improvement(billing): skip the shared usage read while Redis is not ready

Author: waleedlatif1

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

@@ -3,6 +3,8 @@
  * 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 AddressInfo, createServer, type Socket } from 'node:net'
 import type { db } from '@sim/db'
 import * as schema from '@sim/db/schema'
 import { readTestDatabaseUrl, readTestRedisUrl } from '@sim/db/testing/test-infrastructure'
@@ -78,16 +80,14 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
     await connection.end()
   })
 
-  it('serves a sum another process stored without summing the ledger', async () => {
+  it('serves a sum another process stored instead of 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))
@@ -99,7 +99,6 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
     await redis.set(sharedKey(payer), 'not-a-number', 'PX', 30_000)
 
     await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
-    expect(transaction).toHaveBeenCalledTimes(1)
   })
 
   it('never lets a slower, older sum replace one stored while it ran', async () => {
@@ -121,20 +120,45 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
     await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.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,
+  it('sums the ledger without issuing or queuing a command while Redis is not ready', async () => {
+    const notReady = new Redis(redisUrl!, { lazyConnect: true, maxRetriesPerRequest: 0 })
+    redisConfigMockFns.mockGetRedisClient.mockReturnValue(notReady)
+    try {
+      await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
+      expect(notReady.status).toBe('wait')
+
+      await notReady.connect()
+      await notReady.ping()
+      expect(await redis.get(sharedKey(payer))).toBeNull()
+    } finally {
+      notReady.disconnect()
+    }
+  })
+
+  it('sums the ledger promptly when a connected Redis stops answering', async () => {
+    const sockets = new Set<Socket>()
+    /** Completes the client's handshake, then never answers a read. */
+    const silent = createServer((socket) => {
+      sockets.add(socket)
+      socket.on('data', (data) => {
+        const text = data.toString()
+        if (/\bGET\b/i.test(text)) return
+        socket.write('+OK\r\n'.repeat(text.match(/^\*\d+\r\n/gm)?.length ?? 0))
+      })
     })
-    unreachable.on('error', () => undefined)
-    redisConfigMockFns.mockGetRedisClient.mockReturnValue(unreachable)
+    await new Promise<void>((resolve) => silent.listen(0, '127.0.0.1', resolve))
+    const { port } = silent.address() as AddressInfo
+    const hung = new Redis({ host: '127.0.0.1', port, lazyConnect: true, enableReadyCheck: false })
+    redisConfigMockFns.mockGetRedisClient.mockReturnValue(hung)
     try {
+      await hung.connect()
       const startedAt = Date.now()
       await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
       expect(Date.now() - startedAt).toBeLessThan(2_000)
     } finally {
-      unreachable.disconnect()
+      hung.disconnect()
+      for (const socket of sockets) socket.destroy()
+      silent.close()
     }
   })
 })

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

@@ -2,6 +2,7 @@ import { db } from '@sim/db'
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import { randomInt } from '@sim/utils/random'
+import type Redis from 'ioredis'
 import { LRUCache } from 'lru-cache'
 import {
   type BillingEntity,
@@ -37,9 +38,8 @@ export const REPORTING_USAGE_CACHE_TTL_MS = 30_000
 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.
+ * 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
 
@@ -53,14 +53,26 @@ function sharedReportingUsageKey(key: string): string {
   return `usage:reporting:${SHARED_KEY_VERSION}:${key}`
 }
 
+/**
+ * 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.
+ */
+function readyRedisClient(): Redis | null {
+  const redis = getRedisClient()
+  return redis?.status === 'ready' ? redis : null
+}
+
 /**
  * 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.
+ * not ready, 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> {
   try {
-    const redis = getRedisClient()
+    const redis = readyRedisClient()
     if (!redis) return undefined
     const stored = await withinDeadline(
       () => redis.get(sharedReportingUsageKey(key)),
@@ -89,7 +101,7 @@ function warnSharedWriteFailed(error: unknown): void {
  */
 function writeSharedReportingUsageCost(key: string, cost: number): void {
   try {
-    const redis = getRedisClient()
+    const redis = readyRedisClient()
     if (!redis) return
     const ttlMs = REPORTING_USAGE_CACHE_TTL_MS + randomInt(0, SHARED_TTL_JITTER_MS)
     redis

apps/sim/lib/billing/core/usage-threshold-email.integration.ts

@@ -1,13 +1,13 @@
 /**
  * The level-triggered usage threshold email against real claim state in a disposable PostgreSQL
- * schema. Only delivery is stubbed: the mailer is the external boundary, and a sent email is the
- * outcome under test.
+ * schema. Only delivery is stubbed: the mailer is the external boundary, so the outcomes under
+ * test are the messages handed to it and the claim state left in the database.
  */
 import type { db } from '@sim/db'
 import * as schema from '@sim/db/schema'
 import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure'
 import { emailMailerMock, emailMailerMockFns } from '@sim/testing/mocks/email-mailer.mock'
-import { emailTemplatesMock } from '@sim/testing/mocks/email-templates.mock'
+import { emailTemplatesMock, emailTemplatesMockFns } from '@sim/testing/mocks/email-templates.mock'
 import { envFlagsMock, resetEnvFlagsMock, setEnvFlags } from '@sim/testing/mocks/env-flags.mock'
 import { generateId } from '@sim/utils/id'
 import { drizzle } from 'drizzle-orm/postgres-js'
@@ -59,6 +59,40 @@ function notify(
   })
 }
 
+/** Every message handed to the mailer so far, as who received which email. */
+function delivered(): { to: string; subject: string }[] {
+  return mockSendEmail.mock.calls.map(([message]) => ({ to: message.to, subject: message.subject }))
+}
+
+function warning(): { to: string; subject: string } {
+  return {
+    to: `${adminId}@example.com`,
+    subject: emailTemplatesMockFns.mockGetEmailSubject('usage-threshold'),
+  }
+}
+
+function reached(): { to: string; subject: string } {
+  return {
+    to: `${adminId}@example.com`,
+    subject: emailTemplatesMockFns.mockGetLimitEmailSubject('credits', 'reached'),
+  }
+}
+
+/** The organization's persisted threshold claims. */
+async function claims(): Promise<Record<string, number>> {
+  const [row] = await connection<{ limit_notifications: Record<string, number> | null }[]>`
+    SELECT limit_notifications FROM organization WHERE id = ${organizationId}`
+  return row.limit_notifications ?? {}
+}
+
+function claimOf(threshold: 80 | 100, periodStart = SEPTEMBER, limitCents = 10_000) {
+  return {
+    credits: threshold,
+    creditsPeriod: Math.floor(periodStart.getTime() / 86_400_000),
+    creditsLimit: limitCents,
+  }
+}
+
 async function setNotificationsEnabled(enabled: boolean): Promise<void> {
   await connection`INSERT INTO settings VALUES (${adminId}, ${adminId}, ${enabled})
     ON CONFLICT (id) DO UPDATE SET billing_usage_notifications_enabled = ${enabled}`
@@ -97,7 +131,8 @@ describe('usage threshold email', () => {
     await Promise.all([notify(85), notify(85), notify(86)])
     await notify(90)
 
-    expect(mockSendEmail).toHaveBeenCalledTimes(1)
+    expect(delivered()).toEqual([warning()])
+    expect(await claims()).toEqual(claimOf(80))
   })
 
   it('still sends the reached email after the warning, but never the warning after it', async () => {
@@ -106,7 +141,8 @@ describe('usage threshold email', () => {
     await notify(100)
     await notify(85)
 
-    expect(mockSendEmail).toHaveBeenCalledTimes(2)
+    expect(delivered()).toEqual([warning(), reached()])
+    expect(await claims()).toEqual(claimOf(100))
   })
 
   it('re-arms both thresholds whenever the billing period changes, even to an earlier one', async () => {
@@ -115,31 +151,36 @@ describe('usage threshold email', () => {
     await notify(100, { periodStart: OCTOBER })
     await notify(85)
 
-    expect(mockSendEmail).toHaveBeenCalledTimes(4)
+    expect(delivered()).toEqual([reached(), warning(), reached(), warning()])
+    expect(await claims()).toEqual(claimOf(80))
   })
 
   it('warns again at a raised limit after the old one was reached', async () => {
     await notify(100)
     await notify(100, { limit: 125 })
     await notify(110, { limit: 125 })
 
-    expect(mockSendEmail).toHaveBeenCalledTimes(2)
+    expect(delivered()).toEqual([reached(), warning()])
+    expect(await claims()).toEqual(claimOf(80, SEPTEMBER, 12_500))
   })
 
   it('keeps the claim for a later completion when nobody can be notified', async () => {
     await setNotificationsEnabled(false)
     await notify(90)
+    expect(delivered()).toEqual([])
+    expect(await claims()).toEqual({})
+
     await setNotificationsEnabled(true)
     await notify(90)
-
-    expect(mockSendEmail).toHaveBeenCalledTimes(1)
+    expect(delivered()).toEqual([warning()])
   })
 
   it('keeps the claim when a completion recorded no cost', async () => {
     await notify(90, { costDelta: 0 })
-    expect(mockSendEmail).not.toHaveBeenCalled()
+    expect(delivered()).toEqual([])
+    expect(await claims()).toEqual({})
 
     await notify(90)
-    expect(mockSendEmail).toHaveBeenCalledTimes(1)
+    expect(delivered()).toEqual([warning()])
   })
 })

apps/sim/lib/billing/core/usage.test.ts

@@ -331,20 +331,25 @@ describe('maybeSendUsageThresholdEmail', () => {
     resetEnvFlagsMock()
   })
 
+  /** Who received which email, identified by its subject and rendered template. */
+  function sentMessages() {
+    return mockSendEmail.mock.calls.map(([message]) => ({
+      to: message.to,
+      subject: message.subject,
+      html: message.html,
+    }))
+  }
+
   it('emails a paid personal account at 100% with the raise-your-limit template', async () => {
     await maybeSendUsageThresholdEmail({
       ...paidUser,
       usageBefore: 19,
       costDelta: 1,
     })
 
-    expect(mockRenderUsageLimitReached).toHaveBeenCalledWith(
-      expect.objectContaining({ scope: 'user', planName: 'Pro' })
-    )
-    expect(mockRenderCreditsExhausted).not.toHaveBeenCalled()
-    expect(mockRenderUsageThreshold).not.toHaveBeenCalled()
-    expect(mockGetLimitEmailSubject).toHaveBeenCalledWith('credits', 'reached')
-    expect(mockSendEmail).toHaveBeenCalledTimes(1)
+    expect(sentMessages()).toEqual([
+      { to: 'user-1@example.com', subject: 'Limit subject', html: '<html>reached</html>' },
+    ])
   })
 
   it('fans out to org admins at 100% and skips non-admin members', async () => {
@@ -365,12 +370,8 @@ describe('maybeSendUsageThresholdEmail', () => {
       limit: 500,
     })
 
-    expect(mockSendEmail).toHaveBeenCalledTimes(1)
-    expect(mockSendEmail).toHaveBeenCalledWith(
-      expect.objectContaining({ to: 'admin@example.com', emailType: 'notifications' })
-    )
-    expect(mockRenderUsageLimitReached).toHaveBeenCalledWith(
-      expect.objectContaining({ scope: 'organization' })
-    )
+    expect(sentMessages()).toEqual([
+      { to: 'admin@example.com', subject: 'Limit subject', html: '<html>reached</html>' },
+    ])
   })
 })