improvement(billing): share the enterprise reporting usage read across processes

Author: waleedlatif1

apps/sim/lib/billing/core/limit-notifications.ts

@@ -17,6 +17,9 @@ const logger = createLogger('LimitNotifications')
 /** Limit categories that send per-category threshold emails (credits has its own path). */
 export type LimitCategory = Extract<UpgradeReason, 'storage' | 'tables' | 'seats'>
 
+/** Every category whose emailed threshold is persisted, including credits. */
+type ClaimCategory = LimitCategory | Extract<UpgradeReason, 'credits'>
+
 const WARN_THRESHOLD = 80
 const REACH_THRESHOLD = 100
 /** Usage must drop below this band before the same threshold can re-notify (hysteresis). */
@@ -41,7 +44,7 @@ function thresholdFor(percent: number): 0 | 80 | 100 {
 async function claimThreshold(
   scope: 'user' | 'organization',
   id: string,
-  category: LimitCategory,
+  category: ClaimCategory,
   threshold: number
 ): Promise<boolean> {
   const setExpr = sql`jsonb_set(coalesce(${scope === 'user' ? userStats.limitNotifications : organization.limitNotifications}, '{}'::jsonb), ARRAY[${category}], to_jsonb(${threshold}::int))`
@@ -66,6 +69,24 @@ async function claimThreshold(
   return claimed.length > 0
 }
 
+const DAY_MS = 24 * 60 * 60 * 1000
+
+/**
+ * Claim a credits threshold (80 or 100) once per billing period, returning whether THIS call won
+ * it. The stored value is the period's start day followed by the threshold, so it only grows: a
+ * later period outranks every claim of an earlier one and re-arms both thresholds with no reset
+ * write, while within a period a claim of 100 also retires 80, and never the reverse.
+ */
+export function claimCreditsThreshold(
+  scope: 'user' | 'organization',
+  id: string,
+  periodStart: Date,
+  threshold: 80 | 100
+): Promise<boolean> {
+  const periodDay = Math.floor(periodStart.getTime() / DAY_MS)
+  return claimThreshold(scope, id, 'credits', periodDay * 1000 + threshold)
+}
+
 /** Re-arm a category (reset its stored threshold to 0) once usage falls back into the low band. */
 async function rearmThreshold(
   scope: 'user' | 'organization',

apps/sim/lib/billing/core/reporting-usage-cache.integration.ts

@@ -0,0 +1,121 @@
+/**
+ * The shared reporting-usage read against a real ledger in a disposable PostgreSQL schema and a
+ * real Redis. Skipped without `TEST_REDIS_URL`. Each test uses a fresh payer, so the in-process
+ * cache is always cold and every read models a new process.
+ */
+import type { db } from '@sim/db'
+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 { 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 databaseUrl = readTestDatabaseUrl()
+const redisUrl = readTestRedisUrl()
+
+vi.mock('@sim/db', () => ({ db: { transaction }, dbReplica: {} }))
+vi.mock('@/lib/core/config/redis', () => redisConfigMock)
+
+import { readSoftGateUsageCost } from '@/lib/billing/core/reporting-usage-cache'
+import type { BillingEntity, UsageQueryPeriod } from '@/lib/billing/core/usage-log'
+
+const schemaName = `reporting_usage_${generateId().replaceAll('-', '')}`
+const connection = postgres(databaseUrl, {
+  max: 2,
+  prepare: false,
+  connection: { search_path: schemaName },
+  onnotice: () => undefined,
+})
+const database = drizzle(connection, { schema }) as typeof db
+
+const REPORTING: UsageQueryPeriod = {
+  start: new Date('2026-01-01T00:00:00.000Z'),
+  end: new Date('2027-01-01T00:00:00.000Z'),
+  source: 'reporting',
+}
+
+function sharedKey(payer: BillingEntity): string {
+  return `usage:reporting:v1:${payer.type}:${payer.id}:reporting:${REPORTING.start.toISOString()}:${REPORTING.end.toISOString()}`
+}
+
+describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
+  let redis: Redis
+  let payer: BillingEntity
+
+  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,
+      billing_entity_id text, billing_period_start timestamp, billing_period_end timestamp,
+      created_at timestamp NOT NULL
+    )`)
+    transaction.mockImplementation((callback) => database.transaction(callback))
+  })
+
+  beforeEach(async () => {
+    transaction.mockClear()
+    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)
+      VALUES (${generateId()}, 4.25, 'organization', ${payer.id}, '2026-03-01'),
+        (${generateId()}, 1.5, 'organization', ${payer.id}, '2026-06-01'),
+        (${generateId()}, 99, 'organization', ${payer.id}, '2025-12-31')`
+  })
+
+  afterEach(async () => {
+    await redis.del(sharedKey(payer))
+  })
+
+  afterAll(async () => {
+    await redis?.quit()
+    await connection.unsafe(`DROP SCHEMA IF EXISTS "${schemaName}" CASCADE`)
+    await connection.end()
+  })
+
+  it('serves a sum another process stored without summing the ledger', async () => {
+    await redis.set(sharedKey(payer), '12.5', 'PX', 30_000)
+
+    await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(12.5)
+    expect(transaction).not.toHaveBeenCalled()
+  })
+
+  it('sums the ledger exactly on a miss and stores the sum for other processes', async () => {
+    await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
+    expect(transaction).toHaveBeenCalledTimes(1)
+
+    await vi.waitFor(async () => expect(await redis.get(sharedKey(payer))).toBe('5.75'))
+    const ttl = await redis.pttl(sharedKey(payer))
+    expect(ttl).toBeGreaterThan(25_000)
+    expect(ttl).toBeLessThanOrEqual(35_000)
+  })
+
+  it('treats an unreadable stored sum as a miss and overwrites it', async () => {
+    await redis.set(sharedKey(payer), 'not-a-number', 'PX', 30_000)
+
+    await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
+    await vi.waitFor(async () => expect(await redis.get(sharedKey(payer))).toBe('5.75'))
+  })
+
+  it('sums the ledger promptly when Redis is unreachable', async () => {
+    const unreachable = new Redis('redis://127.0.0.1:1', {
+      enableOfflineQueue: true,
+      maxRetriesPerRequest: null,
+      retryStrategy: () => 1_000,
+    })
+    unreachable.on('error', () => undefined)
+    redisConfigMockFns.mockGetRedisClient.mockReturnValue(unreachable)
+    try {
+      const startedAt = Date.now()
+      await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
+      expect(Date.now() - startedAt).toBeLessThan(2_000)
+    } finally {
+      unreachable.disconnect()
+    }
+  })
+})

apps/sim/lib/billing/core/reporting-usage-cache.ts

@@ -1,12 +1,19 @@
 import { db } from '@sim/db'
+import { createLogger } from '@sim/logger'
+import { getErrorMessage } from '@sim/utils/errors'
+import { randomInt } from '@sim/utils/random'
 import { LRUCache } from 'lru-cache'
 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'
 
+const logger = createLogger('ReportingUsageCache')
+
 /**
  * How long a reporting-window usage sum is served before it is summed again.
  *
@@ -18,23 +25,97 @@ import type { DbClient } from '@/lib/db/types'
  * an admission gate lets a payer run on for at most this long past their limit, and a sum at or
  * above the limit is a refusal the true sum would also give. Thirty seconds keeps that overrun
  * small against a year-long allowance while turning a per-event scan into one per window.
+ *
+ * A sum is held both in Redis, shared by every process, and in each process that reads it, so a
+ * served sum can be up to twice this old (plus the Redis expiry's jitter).
  */
 export const REPORTING_USAGE_CACHE_TTL_MS = 30_000
 
+/** Redis expiry is jittered by up to this much, so payers summed together do not expire together. */
+const SHARED_TTL_JITTER_MS = 5_000
+
+/**
+ * How long a read waits on Redis before summing the ledger instead. The shared client queues
+ * commands while disconnected and has long timeouts, so without this a Redis outage would stall
+ * every gate behind it rather than cost one sum.
+ */
+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' }
 
+function sharedReportingUsageKey(key: string): string {
+  return `usage:reporting:${SHARED_KEY_VERSION}:${key}`
+}
+
+/**
+ * A sum another process stored, or `undefined` when there is none to use. Redis being absent,
+ * slow, or failing, and a value that is not a non-negative number, are all misses: the caller
+ * sums the ledger, so the cache can cost a read its latency but never its answer.
+ */
+async function readSharedReportingUsageCost(key: string): Promise<number | undefined> {
+  const redis = getRedisClient()
+  if (!redis) return undefined
+  try {
+    const stored = await withinDeadline(
+      () => redis.get(sharedReportingUsageKey(key)),
+      Date.now() + SHARED_READ_TIMEOUT_MS
+    )
+    if (stored === null) return undefined
+    const cost = Number(stored)
+    if (stored.trim() !== '' && Number.isFinite(cost) && cost >= 0) return cost
+    logger.warn('Discarding unreadable shared reporting usage', { key })
+  } catch (error) {
+    logger.warn('Shared reporting usage read failed; summing the ledger', {
+      error: getErrorMessage(error),
+    })
+  }
+  return undefined
+}
+
+/** Fire-and-forget: a read never waits on, or fails because of, the shared write. */
+function writeSharedReportingUsageCost(key: string, cost: number): void {
+  const redis = getRedisClient()
+  if (!redis) return
+  const ttlMs = REPORTING_USAGE_CACHE_TTL_MS + randomInt(0, SHARED_TTL_JITTER_MS)
+  redis.set(sharedReportingUsageKey(key), String(cost), 'PX', ttlMs).catch((error: unknown) => {
+    logger.warn('Shared reporting usage write failed', { error: getErrorMessage(error) })
+  })
+}
+
+/**
+ * The sum from Redis when another process stored one, else the ledger's exact sum, stored for
+ * the others. Trigger.dev runs each task in a fresh process, so the in-process cache alone is
+ * always cold there; the shared value is what spares those runs the scan.
+ */
+async function sumReportingUsageCost(
+  key: string,
+  entity: BillingEntity,
+  period: ReportingQueryPeriod
+): Promise<number> {
+  const shared = await readSharedReportingUsageCost(key)
+  if (shared !== undefined) return shared
+  const cost = await getBillingPeriodUsageCost(entity, period)
+  writeSharedReportingUsageCost(key, cost)
+  return cost
+}
+
 /**
- * Sums shared across callers, one per payer and window. Every key is an enterprise payer's
- * current window, a few dozen bytes each, so the ceiling sits far above any process's working
- * set and only backstops memory; an eviction inside the TTL costs one extra sum.
+ * 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
+ * above any process's working set and only backstops memory; an eviction inside the TTL costs
+ * one extra read.
  *
- * `fetchMethod` coalesces concurrent misses onto one sum. A rejected sum is evicted rather than
+ * `fetchMethod` coalesces concurrent misses onto one read. A rejected sum is evicted rather than
  * 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 either — usage is written by execution
- * workers in other processes, so the TTL is the real bound.
+ * 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.
  */
 const reportingUsageCache = new LRUCache<
   string,
@@ -43,8 +124,8 @@ const reportingUsageCache = new LRUCache<
 >({
   max: 1_000,
   ttl: REPORTING_USAGE_CACHE_TTL_MS,
-  fetchMethod: (_key, _stale, { context }) =>
-    getBillingPeriodUsageCost(context.entity, context.period),
+  fetchMethod: (key, _stale, { context }) =>
+    sumReportingUsageCost(key, context.entity, context.period),
 })
 
 /**
@@ -72,7 +153,7 @@ 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 for up to {@link REPORTING_USAGE_CACHE_TTL_MS}, since their
+ * served from the shared cache for up to twice {@link REPORTING_USAGE_CACHE_TTL_MS}, 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.
  * Never use it for invoicing, cycle close, an edge-triggered decision, or a read that must see its