diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts b/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts index a8986e95741..e83bed744cf 100644 --- a/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts +++ b/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts @@ -10,27 +10,47 @@ import * as schema from '@sim/db/schema' import { readTestDatabaseUrl, readTestRedisUrl } from '@sim/db/testing/test-infrastructure' import { redisConfigMock, redisConfigMockFns } from '@sim/testing/mocks/redis-config.mock' import { generateId } from '@sim/utils/id' +import { randomInt, randomItem } from '@sim/utils/random' import { drizzle } from 'drizzle-orm/postgres-js' import Redis from 'ioredis' import postgres from 'postgres' import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { transaction } = vi.hoisted(() => ({ transaction: vi.fn() })) +const { transaction, summedPeriods } = vi.hoisted(() => ({ + transaction: vi.fn(), + summedPeriods: [] as UsageQueryPeriod[], +})) const databaseUrl = readTestDatabaseUrl() const redisUrl = readTestRedisUrl() vi.mock('@sim/db', () => ({ db: { transaction }, dbReplica: {} })) vi.mock('@/lib/core/config/redis', () => redisConfigMock) +vi.mock('@/lib/billing/core/usage-log', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + getBillingPeriodUsageCost: (...args: Parameters) => { + summedPeriods.push(args[1]) + return actual.getBillingPeriodUsageCost(...args) + }, + } +}) import { REPORTING_USAGE_CACHE_TTL_MS, readSoftGateUsageCost, } from '@/lib/billing/core/reporting-usage-cache' -import type { BillingEntity, UsageQueryPeriod } from '@/lib/billing/core/usage-log' +import { invalidateReportingUsagePrefixes } from '@/lib/billing/core/reporting-usage-epoch' +import { + type BillingEntity, + type RecordCumulativeUsageParams, + recordCumulativeUsage, + type UsageQueryPeriod, +} from '@/lib/billing/core/usage-log' const schemaName = `reporting_usage_${generateId().replaceAll('-', '')}` const connection = postgres(databaseUrl, { - max: 2, + max: 4, prepare: false, connection: { search_path: schemaName }, onnotice: () => undefined, @@ -43,28 +63,128 @@ const REPORTING: UsageQueryPeriod = { source: 'reporting', } -function sharedKey(payer: BillingEntity): string { - return `usage:reporting:v1:${payer.type}:${payer.id}:reporting:${REPORTING.start.toISOString()}:${REPORTING.end.toISOString()}` +const DAY_MS = 86_400_000 +const MINUTE_MS = 60_000 + +/** A year-long window that started 100 days ago, so the clock is always inside it. */ +function openReportingWindow(): UsageQueryPeriod { + const start = Math.floor((Date.now() - 100 * DAY_MS) / DAY_MS) * DAY_MS + return { start: new Date(start), end: new Date(start + 365 * DAY_MS), source: 'reporting' } +} + +function windowKey(payer: BillingEntity, period: UsageQueryPeriod): string { + return `${payer.type}:${payer.id}:reporting:${period.start.toISOString()}:${period.end.toISOString()}` +} + +function sharedKey(payer: BillingEntity, period: UsageQueryPeriod = REPORTING): string { + return `usage:reporting:v1:${windowKey(payer, period)}` +} + +function prefixKey(payer: BillingEntity, period: UsageQueryPeriod): string { + return `usage:reporting:v1:prefix:${windowKey(payer, period)}` +} + +function prefixLeaseKey(payer: BillingEntity, period: UsageQueryPeriod): string { + return `usage:reporting:v1:prefix-lease:${windowKey(payer, period)}` +} + +function payerEpochKey(payer: BillingEntity): string { + return `usage:reporting:v1:prefix-epoch:${payer.type}:${payer.id}` +} + +type ReportingUsageCacheModule = typeof import('@/lib/billing/core/reporting-usage-cache') + +/** Fresh module instances, so each in-process cache is cold, as in a new process. */ +async function newProcesses(count: number): Promise { + const processes: ReportingUsageCacheModule[] = [] + for (let index = 0; index < count; index++) { + vi.resetModules() + processes.push(await import('@/lib/billing/core/reporting-usage-cache')) + } + return processes +} + +async function readInNewProcess(payer: BillingEntity, period: UsageQueryPeriod): Promise { + const [cache] = await newProcesses(1) + return cache.readSoftGateUsageCost(payer, period) +} + +/** Sums that start at the window's start: a whole-window sum or a prefix, never a tail. */ +function windowStartSums(period: UsageQueryPeriod): UsageQueryPeriod[] { + return summedPeriods.filter((summed) => summed.start.getTime() === period.start.getTime()) +} + +async function insertUsage(payer: BillingEntity, cost: number, createdAt: Date): Promise { + const id = generateId() + await connection`INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at) + VALUES (${id}, ${cost}, ${payer.type}, ${payer.id}, ${createdAt.toISOString()})` + return id +} + +function cumulativeCharge( + payer: BillingEntity, + eventKey: string, + cost: number +): RecordCumulativeUsageParams { + return { + userId: 'actor', + workspaceId: 'workspace', + billingEntity: payer, + billingPeriod: { start: REPORTING.start, end: REPORTING.end }, + source: 'workspace-chat', + model: 'test-model', + eventKey, + cost, + } +} + +/** A cumulative charge whose row was first written at `createdAt`. */ +async function startCumulativeCharge( + payer: BillingEntity, + cost: number, + createdAt: Date +): Promise { + const eventKey = `update-cost:${generateId()}` + await recordCumulativeUsage(cumulativeCharge(payer, eventKey, cost)) + await connection`UPDATE usage_log SET created_at = ${createdAt.toISOString()} + WHERE event_key = ${eventKey}` + return eventKey +} + +async function exactSum(payer: BillingEntity, period: UsageQueryPeriod): Promise { + const [row] = await connection<{ cost: string }[]>`SELECT COALESCE(SUM(cost), 0) AS cost + FROM usage_log WHERE billing_entity_type = ${payer.type} AND billing_entity_id = ${payer.id} + AND created_at >= ${period.start.toISOString()} AND created_at < ${period.end.toISOString()}` + return Number(row?.cost ?? 0) } describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => { let redis: Redis let payer: BillingEntity + async function prefixStored(payer: BillingEntity, period: UsageQueryPeriod): Promise { + await vi.waitFor(async () => expect(await redis.exists(prefixKey(payer, period))).toBe(1)) + } + beforeAll(async () => { redis = new Redis(redisUrl!, { lazyConnect: true, maxRetriesPerRequest: 0 }) await redis.connect() await connection.unsafe(`CREATE SCHEMA "${schemaName}"`) await connection.unsafe(`CREATE TABLE usage_log ( - id text PRIMARY KEY, cost numeric NOT NULL, billing_entity_type text, + id text PRIMARY KEY, user_id text, category text, source text, description text, + metadata jsonb, cost numeric NOT NULL, event_key text, billing_entity_type text, billing_entity_id text, billing_period_start timestamp, billing_period_end timestamp, - created_at timestamp NOT NULL - )`) + workspace_id text, workflow_id text, execution_id text, + created_at timestamp NOT NULL DEFAULT now() + ); + CREATE UNIQUE INDEX usage_log_event_key_unique ON usage_log(event_key) + WHERE event_key IS NOT NULL`) transaction.mockImplementation((callback) => database.transaction(callback)) }) beforeEach(async () => { transaction.mockClear() + summedPeriods.length = 0 redisConfigMockFns.mockGetRedisClient.mockReturnValue(redis) payer = { type: 'organization', id: generateId() } await connection`INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at) @@ -75,7 +195,15 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => { afterEach(async () => { vi.useRealTimers() - await redis.del(sharedKey(payer)) + await redis.del( + ...[REPORTING, openReportingWindow()].flatMap((period) => [ + sharedKey(payer, period), + prefixKey(payer, period), + prefixLeaseKey(payer, period), + ]), + payerEpochKey(payer), + 'usage:reporting:v1:prefix-epoch' + ) }) afterAll(async () => { @@ -149,6 +277,170 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => { expect(await redis.get(sharedKey(payer))).toBe('9.99') }) + it('sums the whole window in one query when Redis is absent', async () => { + redisConfigMockFns.mockGetRedisClient.mockReturnValue(null) + const period = openReportingWindow() + await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS)) + await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS)) + + await expect(readInNewProcess(payer, period)).resolves.toBe(5) + expect(summedPeriods).toEqual([period]) + }) + + it('leases a due refresh to one of the processes missing together', async () => { + const period = openReportingWindow() + await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS)) + await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS)) + await expect(readInNewProcess(payer, period)).resolves.toBe(5) + await insertUsage(payer, 4, new Date(period.start.getTime() + 2 * DAY_MS)) + + const now = Date.now() + vi.useFakeTimers({ toFake: ['Date'] }) + vi.setSystemTime(now + 4 * MINUTE_MS) + await redis.del(sharedKey(payer, period), prefixLeaseKey(payer, period)) + const processes = await newProcesses(5) + const costs = await Promise.all( + processes.map((cache) => cache.readSoftGateUsageCost(payer, period)) + ) + + expect(windowStartSums(period)).toHaveLength(2) + expect(costs).toContain(9) + for (const cost of costs) expect([5, 9]).toContain(cost) + }) + + it('sums the whole window at once when no prefix is stored and another process is summing one', async () => { + const period = openReportingWindow() + await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS)) + await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS)) + await redis.set(prefixLeaseKey(payer, period), '1', 'PX', MINUTE_MS) + const [cache] = await newProcesses(1) + + const startedAt = performance.now() + await expect(cache.readSoftGateUsageCost(payer, period)).resolves.toBe(5) + expect(performance.now() - startedAt).toBeLessThan(1_000) + expect(summedPeriods).toEqual([period]) + }) + + it('never reads above the ledger as it grows, and matches it once the prefix is refreshed', async () => { + const period = openReportingWindow() + const oldRows = [ + await insertUsage(payer, 1.25, new Date(period.start.getTime() + DAY_MS)), + await insertUsage(payer, 0.5, new Date(Date.now() - 10 * DAY_MS)), + ] + await expect(readInNewProcess(payer, period)).resolves.toBe(1.75) + + for (let round = 0; round < 8; round++) { + const recent = new Date(Date.now() - randomInt(0, 4 * MINUTE_MS)) + await insertUsage(payer, randomInt(1, 500) / 100, recent) + if (round % 2 === 0) { + const late = new Date(period.start.getTime() + randomInt(1, 90) * DAY_MS) + await insertUsage(payer, randomInt(1, 500) / 100, late) + } else { + await connection`UPDATE usage_log SET cost = cost + ${randomInt(1, 500) / 100} + WHERE id = ${randomItem(oldRows)}` + } + await redis.del(sharedKey(payer, period)) + + const served = await readInNewProcess(payer, period) + expect(served).toBeLessThanOrEqual((await exactSum(payer, period)) + 1e-9) + } + expect(windowStartSums(period)).toHaveLength(1) + + const now = Date.now() + vi.useFakeTimers({ toFake: ['Date'] }) + vi.setSystemTime(now + 4 * MINUTE_MS) + await redis.del(sharedKey(payer, period), prefixLeaseKey(payer, period)) + + const refreshed = await readInNewProcess(payer, period) + expect(refreshed).toBeCloseTo(await exactSum(payer, period), 9) + expect(windowStartSums(period)).toHaveLength(2) + }) + + it('never reads above the ledger after ledger rows are deleted and prefixes are invalidated', async () => { + const period = openReportingWindow() + await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS)) + const departedRow = await insertUsage(payer, 6, new Date(period.start.getTime() + 2 * DAY_MS)) + await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS)) + await expect(readInNewProcess(payer, period)).resolves.toBe(11) + await prefixStored(payer, period) + + await connection`DELETE FROM usage_log WHERE id = ${departedRow}` + await invalidateReportingUsagePrefixes() + await redis.del(sharedKey(payer, period)) + + await expect(readInNewProcess(payer, period)).resolves.toBe(await exactSum(payer, period)) + }) + + it('sees a top-up to a row inside the stored prefix on the next read', async () => { + const period = openReportingWindow() + await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS)) + const eventKey = await startCumulativeCharge( + payer, + 2, + new Date(period.start.getTime() + DAY_MS) + ) + await expect(readInNewProcess(payer, period)).resolves.toBe(5) + await prefixStored(payer, period) + + await recordCumulativeUsage(cumulativeCharge(payer, eventKey, 6)) + await vi.waitFor(async () => expect(await redis.exists(payerEpochKey(payer))).toBe(1)) + await redis.del(sharedKey(payer, period)) + + await expect(readInNewProcess(payer, period)).resolves.toBe(9) + expect(await exactSum(payer, period)).toBe(9) + }) + + it('leaves the prefix in place for a top-up to a row still in the tail, and reads it exactly', async () => { + const period = openReportingWindow() + await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS)) + const eventKey = await startCumulativeCharge(payer, 3, new Date(Date.now() - MINUTE_MS)) + await expect(readInNewProcess(payer, period)).resolves.toBe(5) + await prefixStored(payer, period) + + await recordCumulativeUsage(cumulativeCharge(payer, eventKey, 7)) + await redis.del(sharedKey(payer, period)) + + await expect(readInNewProcess(payer, period)).resolves.toBe(9) + expect(await redis.exists(payerEpochKey(payer))).toBe(0) + expect(windowStartSums(period)).toHaveLength(1) + }) + + it('never over-counts a top-up that commits while a prefix refresh sums', async () => { + const period = openReportingWindow() + await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS)) + const eventKey = await startCumulativeCharge( + payer, + 2, + new Date(period.start.getTime() + DAY_MS) + ) + const [refreshing] = await newProcesses(1) + transaction.mockImplementationOnce(async (callback) => { + const result = await database.transaction(callback) + await recordCumulativeUsage(cumulativeCharge(payer, eventKey, 6)) + return result + }) + + const served = await refreshing.readSoftGateUsageCost(payer, period) + expect(served).toBeLessThanOrEqual(await exactSum(payer, period)) + await prefixStored(payer, period) + await vi.waitFor(async () => expect(await redis.exists(payerEpochKey(payer))).toBe(1)) + await redis.del(sharedKey(payer, period)) + + await expect(readInNewProcess(payer, period)).resolves.toBe(await exactSum(payer, period)) + }) + + it('serves a stale prefix while another process holds the refresh lease', async () => { + const period = openReportingWindow() + await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS)) + await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS)) + const staleCutoff = Date.now() - 10 * MINUTE_MS + await redis.set(prefixKey(payer, period), `::${staleCutoff}:7`, 'PX', MINUTE_MS) + await redis.set(prefixLeaseKey(payer, period), '1', 'PX', MINUTE_MS) + + await expect(readInNewProcess(payer, period)).resolves.toBe(10) + expect(windowStartSums(period)).toHaveLength(0) + }) + it('sums the ledger when the Redis client cannot be built', async () => { redisConfigMockFns.mockGetRedisClient.mockImplementation(() => { throw new Error('Invalid Redis configuration') diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.ts b/apps/sim/lib/billing/core/reporting-usage-cache.ts index 56ac65d00d0..ba0ecf7c30a 100644 --- a/apps/sim/lib/billing/core/reporting-usage-cache.ts +++ b/apps/sim/lib/billing/core/reporting-usage-cache.ts @@ -4,12 +4,20 @@ import { getErrorMessage } from '@sim/utils/errors' import { randomInt } from '@sim/utils/random' import type Redis from 'ioredis' import { LRUCache } from 'lru-cache' +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' +import { + PREFIX_CUTOFF_LAG_MS, + payerPrefixEpochKey, + readyRedisClient, + SHARED_KEY_VERSION, + SHARED_PREFIX_EPOCH_KEY, + SHARED_READ_TIMEOUT_MS, +} from '@/lib/billing/core/reporting-usage-epoch' import { type BillingEntity, getBillingPeriodUsageCost, type UsageQueryPeriod, } from '@/lib/billing/core/usage-log' -import { getRedisClient } from '@/lib/core/config/redis' import { withinDeadline } from '@/lib/core/utils/deadline' import type { DbClient } from '@/lib/db/types' @@ -32,6 +40,9 @@ const logger = createLogger('ReportingUsageCache') * moment, so it is served from Redis for at most max(this TTL + jitter, the sum's duration + * {@link MIN_SHARED_TTL_MS}) after it began (plus any reconnect delay for a resent write), plus * up to this TTL again in the reading process. + * + * While Redis is ready, a sum is built from a shared {@link LedgerPrefix} plus the rows since its + * cutoff, which adds the prefix's own lag on top of this one; see {@link PREFIX_REFRESH_AFTER_MS}. */ export const REPORTING_USAGE_CACHE_TTL_MS = 30_000 @@ -46,15 +57,6 @@ const SHARED_TTL_JITTER_MS = 5_000 */ const MIN_SHARED_TTL_MS = 5_000 -/** - * How long a read waits on a connected Redis before summing the ledger instead, so a socket that - * has silently stopped answering costs one sum rather than the shared client's long timeouts. - */ -const SHARED_READ_TIMEOUT_MS = 250 - -/** Bump when a stored sum's meaning changes; old entries are then ignored. */ -const SHARED_KEY_VERSION = 'v1' - /** A usage window known to be an enterprise reporting window — the only kind this cache serves. */ type ReportingQueryPeriod = UsageQueryPeriod & { source: 'reporting' } @@ -63,15 +65,65 @@ function sharedReportingUsageKey(key: string): string { } /** - * The shared client, only while its connection is ready. A disconnected client parks commands in - * its offline queue and replays them on reconnect, which would land a stale sum with a fresh - * expiry, and `NX` would then keep newer sums out; so while it is not ready, no command is issued - * at all and the ledger answers instead. Trigger.dev's `init` hook warms the connection, so it is - * ready by the time a task reads. + * How long after it was read a prefix is summed again, by one process. Until then a refill sums + * only the rows since the cutoff, so a payer's year-long scan runs about once per this interval + * instead of once per {@link REPORTING_USAGE_CACHE_TTL_MS} in every process that misses. + * + * The prefix freezes the rows before its cutoff, and the ledger only grows (see + * {@link REPORTING_USAGE_CACHE_TTL_MS}), so prefix plus tail never exceeds the true sum. A + * top-up to a row old enough to sit in a prefix — a long-running cumulative charge — discards + * the payer's prefixes once it commits (`invalidatePayerPrefixesForTopUp`), so the next read + * sums the window again and sees it as soon as it would have without the prefix. What a prefix + * still omits until it is read again is a row whose transaction outlived + * {@link PREFIX_CUTOFF_LAG_MS}, so this interval is the extra lag that rare write sees. Deleting + * a user is the one write that lowers the ledger; it discards every prefix through + * `invalidateReportingUsagePrefixes`. */ -function readyRedisClient(): Redis | null { - const redis = getRedisClient() - return redis?.status === 'ready' ? redis : null +const PREFIX_REFRESH_AFTER_MS = 3 * 60_000 + +/** + * The refresh lease is held for as long as one ledger sum may run, so a second process starts + * the same sum only once the first has certainly ended, and a holder that died frees it then. + * Nothing tells a stranded lease from a live one, and nothing needs to: while it is held, every + * other reader serves the stored prefix or, with none, sums the whole window. + */ +const PREFIX_LEASE_MS = USAGE_LEDGER_STATEMENT_TIMEOUT_MS + +/** + * How long after it was read a prefix stays in Redis, served while another process holds the + * refresh lease: the refresh interval plus one full lease. This is the most a prefix can lag + * when a refresh fails; past it a reader sums the window exactly. + */ +const PREFIX_RETAIN_MS = PREFIX_REFRESH_AFTER_MS + PREFIX_LEASE_MS + +/** The global and payer prefix epochs current when a prefix was summed; `''` while unset. */ +interface PrefixEpochs { + epoch: string + payerEpoch: string +} + +/** + * The window's usage before `cutoff`, as one process summed it under the epochs it read first. + * Stored in Redis as `:::`, keyed like the full sum, so it + * only ever serves that payer's same window, and only while both epochs are still current. + */ +interface LedgerPrefix extends PrefixEpochs { + cutoff: number + cost: number +} + +function sharedPrefixKey(key: string): string { + return `usage:reporting:${SHARED_KEY_VERSION}:prefix:${key}` +} + +function sharedPrefixLeaseKey(key: string): string { + return `usage:reporting:${SHARED_KEY_VERSION}:prefix-lease:${key}` +} + +/** A stored cost, or `undefined` unless it is a non-negative number. */ +function parseStoredCost(text: string): number | undefined { + const cost = Number(text) + return text.trim() !== '' && Number.isFinite(cost) && cost >= 0 ? cost : undefined } /** @@ -88,8 +140,8 @@ async function readSharedReportingUsageCost(key: string): Promise= 0) return cost + const cost = parseStoredCost(stored) + if (cost !== undefined) return cost logger.warn('Discarding unreadable shared reporting usage', { key }) } catch (error) { logger.warn('Shared reporting usage read failed; summing the ledger', { @@ -144,11 +196,157 @@ async function sumReportingUsageCost( const shared = await readSharedReportingUsageCost(key) if (shared !== undefined) return shared const sumStartedAt = Date.now() - const cost = await getBillingPeriodUsageCost(entity, period) + const prefix = await resolveLedgerPrefix(key, entity, period) + const [prefixCost, tailCost] = await Promise.all([ + prefix?.cost ?? 0, + getBillingPeriodUsageCost( + entity, + prefix ? { ...period, start: new Date(prefix.cutoff) } : period + ), + ]) + const cost = prefixCost + tailCost writeSharedReportingUsageCost(key, cost, sumStartedAt) return cost } +interface StoredPrefixState { + /** The epochs current now, which a prefix summed by this read is stored under. */ + epochs: PrefixEpochs + /** The stored prefix, when there is one for this window under both current epochs. */ + prefix: LedgerPrefix | null +} + +/** + * The stored prefix and the current epochs in one round trip, or `undefined` when Redis did not + * answer in time or failed, which ends the prefix path for this read. A prefix summed under an + * earlier global or payer epoch, or that does not fit this window, is no prefix. + */ +async function readStoredPrefix( + redis: Redis, + key: string, + entity: BillingEntity, + period: ReportingQueryPeriod +): Promise { + try { + const [stored, epoch, payerEpoch] = await withinDeadline( + () => redis.mget(sharedPrefixKey(key), SHARED_PREFIX_EPOCH_KEY, payerPrefixEpochKey(entity)), + Date.now() + SHARED_READ_TIMEOUT_MS + ) + const epochs: PrefixEpochs = { epoch: epoch ?? '', payerEpoch: payerEpoch ?? '' } + if (!stored) return { epochs, prefix: null } + const [storedEpoch, storedPayerEpoch, cutoffText = '', costText = '', ...rest] = + stored.split(':') + const cutoff = Number(cutoffText) + const cost = parseStoredCost(costText) + if ( + rest.length === 0 && + Number.isSafeInteger(cutoff) && + cutoff > period.start.getTime() && + cutoff < period.end.getTime() && + cost !== undefined + ) { + const current = storedEpoch === epochs.epoch && storedPayerEpoch === epochs.payerEpoch + return { epochs, prefix: current ? { ...epochs, cutoff, cost } : null } + } + logger.warn('Discarding unreadable shared reporting usage prefix', { key }) + return { epochs, prefix: null } + } catch (error) { + logger.warn('Shared reporting usage prefix read failed; summing the ledger', { + error: getErrorMessage(error), + }) + return undefined + } +} + +/** + * Whether this process took the refresh lease. It is never released: the stored prefix ends the + * need for it, and expiry frees it for the next refresh. A lease whose `SET` failed or timed out + * may still have landed, held by nobody, for {@link PREFIX_LEASE_MS}. Not `acquireLock`, which + * issues on a client that is not ready and has no read deadline. + */ +async function acquirePrefixLease(redis: Redis, key: string): Promise { + try { + const result = await withinDeadline( + () => redis.set(sharedPrefixLeaseKey(key), '1', 'PX', PREFIX_LEASE_MS, 'NX'), + Date.now() + SHARED_READ_TIMEOUT_MS + ) + return result === 'OK' + } catch (error) { + logger.warn('Shared reporting usage prefix lease failed', { + error: getErrorMessage(error), + }) + return false + } +} + +/** + * Fire-and-forget, and only on a ready client, like {@link writeSharedReportingUsageCost}. A + * plain `SET` replaces the stored prefix: every prefix is a lower bound for its own cutoff, so + * two refreshes racing can leave either one, and the tail always resumes from the cutoff stored + * with it. + */ +function writeStoredPrefix(key: string, prefix: LedgerPrefix): void { + try { + const redis = readyRedisClient() + if (!redis) return + redis + .set( + sharedPrefixKey(key), + `${prefix.epoch}:${prefix.payerEpoch}:${prefix.cutoff}:${prefix.cost}`, + 'PX', + prefix.cutoff + PREFIX_CUTOFF_LAG_MS + PREFIX_RETAIN_MS - Date.now() + ) + .catch(warnSharedWriteFailed) + } catch (error) { + warnSharedWriteFailed(error) + } +} + +/** A prefix to add the tail to; `cost` is still pending while this read sums it. */ +interface ResolvedPrefix { + cutoff: number + cost: number | Promise +} + +/** + * The window's usage before a recent cutoff, shared across processes, or `undefined` to sum the + * whole window: Redis is not ready or failing, the window does not span the cutoff, or no + * prefix is stored and another process holds the lease. A fresh stored prefix is served as is. + * Otherwise the process that takes the lease sums a new prefix alongside its tail, and the rest + * serve the stored one while it does. + */ +async function resolveLedgerPrefix( + key: string, + entity: BillingEntity, + period: ReportingQueryPeriod +): Promise { + let redis: Redis | null + try { + redis = readyRedisClient() + } catch { + return undefined + } + if (!redis) return undefined + const cutoff = Date.now() - PREFIX_CUTOFF_LAG_MS + if (cutoff <= period.start.getTime() || cutoff >= period.end.getTime()) return undefined + + const state = await readStoredPrefix(redis, key, entity, period) + if (!state) return undefined + const { epochs, prefix: stored } = state + if (stored && cutoff - stored.cutoff < PREFIX_REFRESH_AFTER_MS) return stored + + if (await acquirePrefixLease(redis, key)) { + const cost = getBillingPeriodUsageCost(entity, { ...period, end: new Date(cutoff) }).then( + (prefixCost) => { + writeStoredPrefix(key, { ...epochs, cutoff, cost: prefixCost }) + return prefixCost + } + ) + return { cutoff, cost } + } + return stored ?? undefined +} + /** * Sums held by this process, one per payer and window, in front of the shared Redis value. Every * key is an enterprise payer's current window, a few dozen bytes each, so the ceiling sits far @@ -159,9 +357,9 @@ async function sumReportingUsageCost( * stored (`noDeleteOnFetchRejection` and `allowStaleOnFetchRejection` stay off), so every caller * of that read sees the error it would have seen uncached and the next call sums again. There is * no settle deadline: the sum runs under the ledger's own `statement_timeout`, so the database - * ends a slow one. There is deliberately no invalidator or lock either — usage is written by - * execution workers in other processes, so the TTL is the real bound, and concurrent misses in - * different processes each sum once. + * ends a slow one. There is deliberately no invalidator either — usage is written by execution + * workers in other processes, so the TTL is the real bound. Concurrent misses in different + * processes each sum the short tail; only the prefix refresh is leased to one of them. */ const reportingUsageCache = new LRUCache< string, @@ -199,10 +397,10 @@ async function readCachedReportingUsageCost( /** * Period usage for a soft reader: an admission check, a display, or a level-triggered * notification that tolerates the cache's bounded under-count. Enterprise reporting windows are - * served from the shared cache, within the lag {@link REPORTING_USAGE_CACHE_TTL_MS} describes, - * since their year-long sum is the expensive one; every other period is summed exactly, as - * before. A read on a caller's own executor (a transaction or a replica) keeps its own snapshot - * and is never shared. + * served from the shared cache, within the lag {@link REPORTING_USAGE_CACHE_TTL_MS} and + * {@link PREFIX_REFRESH_AFTER_MS} describe, since their year-long sum is the expensive one; every + * other period is summed exactly. A read on a caller's own executor (a transaction or a replica) + * keeps its own snapshot and is never shared. * Never use it for invoicing, cycle close, an edge-triggered decision, or a read that must see its * own write. */ diff --git a/apps/sim/lib/billing/core/reporting-usage-epoch.ts b/apps/sim/lib/billing/core/reporting-usage-epoch.ts new file mode 100644 index 00000000000..4188584edae --- /dev/null +++ b/apps/sim/lib/billing/core/reporting-usage-epoch.ts @@ -0,0 +1,106 @@ +import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' +import { generateShortId } from '@sim/utils/id' +import type Redis from 'ioredis' +import type { BillingEntity } from '@/lib/billing/core/usage-log' +import { getRedisClient } from '@/lib/core/config/redis' +import { withinDeadline } from '@/lib/core/utils/deadline' + +const logger = createLogger('ReportingUsageEpoch') + +/** Bump when a stored sum's meaning changes; old entries are then ignored. */ +export const SHARED_KEY_VERSION = 'v1' + +/** + * How long a read waits on a connected Redis before summing the ledger instead, so a socket that + * has silently stopped answering costs one sum rather than the shared client's long timeouts. + */ +export const SHARED_READ_TIMEOUT_MS = 250 + +/** + * How far behind the clock a prefix's cutoff sits. A row's `created_at` is its transaction's + * start on the database clock, so a row stamped just before the cutoff can still be uncommitted + * when the prefix is read; this lag puts every such row in the tail, which is summed on every + * refill, rather than in a prefix that would miss it until its next refresh. It also absorbs + * skew between this process's clock and the database's. + */ +export const PREFIX_CUTOFF_LAG_MS = 5 * 60_000 + +/** + * The age past which a topped-up row may sit inside a stored prefix, so its top-up discards the + * payer's prefixes. A minute short of {@link PREFIX_CUTOFF_LAG_MS}: every cutoff is at least the + * lag behind its reader's clock, and the margin covers skew between that clock, this one, and + * the database's that stamped the row. A younger row is still in every tail. + */ +const TOP_UP_INVALIDATES_AFTER_MS = PREFIX_CUTOFF_LAG_MS - 60_000 + +/** + * How long a bumped prefix epoch is kept. Every prefix written before the bump has expired well + * before this, so the epoch reading as absent again afterwards cannot revive one. + */ +const PREFIX_EPOCH_RETAIN_MS = 60 * 60_000 + +/** + * One value for every payer: a prefix is served only while the epoch it was summed under is + * still current, so replacing the epoch discards every stored prefix at once. + */ +export const SHARED_PREFIX_EPOCH_KEY = `usage:reporting:${SHARED_KEY_VERSION}:prefix-epoch` + +/** Like {@link SHARED_PREFIX_EPOCH_KEY}, for one payer's prefixes only. */ +export function payerPrefixEpochKey(entity: BillingEntity): string { + return `${SHARED_PREFIX_EPOCH_KEY}:${entity.type}:${entity.id}` +} + +/** + * The shared client, only while its connection is ready. A disconnected client parks commands in + * its offline queue and replays them on reconnect, which would land a stale sum with a fresh + * expiry, and `NX` would then keep newer sums out; so while it is not ready, no command is issued + * at all and the ledger answers instead. Trigger.dev's `init` hook warms the connection, so it is + * ready by the time a task reads. + */ +export function readyRedisClient(): Redis | null { + const redis = getRedisClient() + return redis?.status === 'ready' ? redis : null +} + +async function replaceEpoch(epochKey: string): Promise { + const redis = readyRedisClient() + if (!redis) return + await withinDeadline( + () => redis.set(epochKey, generateShortId(), 'PX', PREFIX_EPOCH_RETAIN_MS), + Date.now() + SHARED_READ_TIMEOUT_MS + ) +} + +/** + * Discards every stored prefix, for a write that removes ledger rows — a user's deletion, whose + * rows cascade away. Call it once that write has committed: a refresh that read the old epoch + * stores its prefix under it, where nothing reads it, and one that reads the new epoch sums after + * the deletion. Best effort: without a ready Redis there is no prefix to discard, and a failed + * bump leaves the deleted usage counted until the stored prefixes expire. + */ +export async function invalidateReportingUsagePrefixes(): Promise { + try { + await replaceEpoch(SHARED_PREFIX_EPOCH_KEY) + } catch (error) { + logger.warn('Shared reporting usage prefix invalidation failed', { + error: getErrorMessage(error), + }) + } +} + +/** + * Discards the payer's stored prefixes when a committed top-up raised a row old enough to sit + * inside one, so the next read sums the window again and sees the top-up as soon as an uncached + * read would. Like {@link invalidateReportingUsagePrefixes}, a refresh racing it stores under + * the epoch it read, where nothing reads it once this lands. Fire-and-forget and best effort: a + * failed bump leaves the top-up unseen until the payer's prefixes expire. + */ +export function invalidatePayerPrefixesForTopUp(entity: BillingEntity, rowCreatedAt: Date): void { + if (rowCreatedAt.getTime() >= Date.now() - TOP_UP_INVALIDATES_AFTER_MS) return + replaceEpoch(payerPrefixEpochKey(entity)).catch((error) => { + logger.warn('Shared reporting usage payer prefix invalidation failed', { + error: getErrorMessage(error), + }) + }) +} diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index 204a10f4e24..78bfb51ac06 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -24,6 +24,7 @@ import { resolveSubscriptionUsagePeriod, type UsagePeriodSource, } from '@/lib/billing/core/reporting-period' +import { invalidatePayerPrefixesForTopUp } from '@/lib/billing/core/reporting-usage-epoch' import { apportionCredits } from '@/lib/billing/credits/conversion' import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' import type { InternalUsageLogSource } from '@/lib/billing/usage-sources' @@ -815,6 +816,7 @@ export async function recordCumulativeUsage( const stageDurationsMs: Partial> = {} let succeeded = false let pgCode: string | undefined + let toppedUpRowCreatedAt: Date | undefined const enterStage = (nextStage: CumulativeUsageStage) => { const now = Date.now() stageDurationsMs[stage] = now - stageStartedAt @@ -854,6 +856,7 @@ export async function recordCumulativeUsage( billingEntityId: usageLog.billingEntityId, billingPeriodStart: usageLog.billingPeriodStart, billingPeriodEnd: usageLog.billingPeriodEnd, + createdAt: usageLog.createdAt, }) .from(usageLog) .where( @@ -955,6 +958,7 @@ export async function recordCumulativeUsage( metadata: periodUsageMetadata(metadata, otherRows) ?? null, }) .where(eq(usageLog.id, latest.id)) + toppedUpRowCreatedAt = latest.createdAt enterStage('commit') return { billed: true, delta, total: newTotal, billingPeriod: latestPeriod } } @@ -987,6 +991,9 @@ export async function recordCumulativeUsage( } }) succeeded = true + if (toppedUpRowCreatedAt) { + invalidatePayerPrefixesForTopUp(billingContext.billingEntity, toppedUpRowCreatedAt) + } return result } catch (error) { pgCode = getPostgresErrorCode(error) diff --git a/apps/sim/lib/users/account-deletion.ts b/apps/sim/lib/users/account-deletion.ts index c997c5ebec4..221b42e7153 100644 --- a/apps/sim/lib/users/account-deletion.ts +++ b/apps/sim/lib/users/account-deletion.ts @@ -27,6 +27,7 @@ import type { AccountDeletionResource, } from '@/lib/api/contracts/user' import { getHighestPriorityPersonalSubscription } from '@/lib/billing/core/plan' +import { invalidateReportingUsagePrefixes } from '@/lib/billing/core/reporting-usage-epoch' import { acquireOrganizationMutationLock, isSoleOwnerOfPaidOrganization, @@ -755,7 +756,8 @@ async function announceCancelledTableWork( * rather than treated as failures. * 3. **Purge storage last**, once that transaction has committed. Object * deletion cannot be rolled back, so it must not precede the point of no - * return. + * return. Before it, the account's ledger rows having cascaded away, the + * shared reporting-usage prefixes that may still count them are discarded. * * The ordering inside step 2 is load-bearing too: Postgres evaluates the * `NO ACTION` check on `workspace.billed_account_user_id` *before* the `owner_id` @@ -982,6 +984,8 @@ export async function deleteUserAccount(userId: string): Promise