Skip to content

Commit 10282f2

Browse files
authored
improvement(billing): share the enterprise reporting usage read across processes (#8325)
- Enterprise reporting-window usage sums (a year-long ledger scan) are now shared across processes through Redis with a 30s TTL. Trigger.dev runs every task in a fresh process, so the existing in-process LRU was always cold there, and the scan ran on nearly every document-processing and execution admission check - The in-process LRU stays in front and still coalesces concurrent misses. The Redis GET has a 250ms deadline, and any Redis error, timeout, or unreadable value falls through to the exact sum. The write is a fire-and-forget `SET … NX` with a jittered TTL, so a slower, older sum never overwrites a fresher one or extends its life. There is no lock or lease - The usage threshold email is now level-triggered, instead of being edge-triggered off an exact before/after org sum on every workflow completion. A claim keyed on (billing period, limit) (`claimCreditsThreshold`) sends each threshold at most once. A new period or a changed limit re-arms it with no reset write, and concurrent completions can't both send - Recipients are resolved before claiming, so a period's email isn't used up when nobody can receive it. Zero-cost completions don't claim. The personal usage baseline and the claimed period come from the same billing context - Invoicing, cycle close, overage, and threshold billing still read the ledger exactly - Rollout note: accounts already at or above 80% this billing period get one threshold email after deploy; there is no backfill
1 parent 3945364 commit 10282f2

12 files changed

Lines changed: 759 additions & 198 deletions

File tree

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

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,24 @@
11
import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing/mocks/database.mock'
22
import { emailMailerMock, emailMailerMockFns } from '@sim/testing/mocks/email-mailer.mock'
33
import { emailTemplatesMock, emailTemplatesMockFns } from '@sim/testing/mocks/email-templates.mock'
4+
import {
5+
emailUnsubscribeMock,
6+
emailUnsubscribeMockFns,
7+
} from '@sim/testing/mocks/email-unsubscribe.mock'
48
import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing/mocks/env-flags.mock'
59
import { schemaMock } from '@sim/testing/mocks/schema.mock'
610
import { resetUrlsMock, urlsMockFns } from '@sim/testing/mocks/urls.mock'
711
import { workspaceAuthzMock } from '@sim/testing/mocks/workspace-authz.mock'
812
import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
913

10-
const { getEmailPreferencesMock } = vi.hoisted(() => ({
11-
getEmailPreferencesMock: vi.fn(() => Promise.resolve(null as unknown)),
12-
}))
13-
1414
vi.mock('@/lib/messaging/email/mailer', () => emailMailerMock)
15-
vi.mock('@/lib/messaging/email/unsubscribe', () => ({
16-
getEmailPreferences: getEmailPreferencesMock,
17-
}))
15+
vi.mock('@/lib/messaging/email/unsubscribe', () => emailUnsubscribeMock)
1816
vi.mock('@/components/emails', () => emailTemplatesMock)
1917
vi.mock('@sim/platform-authz/workspace', () => workspaceAuthzMock)
2018

2119
import { maybeSendLimitThresholdEmail } from '@/lib/billing/core/limit-notifications'
2220

21+
const getEmailPreferencesMock = emailUnsubscribeMockFns.mockGetEmailPreferences
2322
const sendEmailSpy = emailMailerMockFns.mockSendEmail
2423
const renderMock = emailTemplatesMockFns.mockRenderLimitThresholdEmail
2524
const subjectMock = emailTemplatesMockFns.mockGetLimitEmailSubject

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

Lines changed: 89 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { db } from '@sim/db'
22
import { member, organization, settings, user, userStats } from '@sim/db/schema'
33
import { createLogger } from '@sim/logger'
44
import { isOrgAdminRole } from '@sim/platform-authz/workspace'
5-
import { and, eq, sql } from 'drizzle-orm'
5+
import { and, eq, type SQL, sql } from 'drizzle-orm'
66
import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
77
import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
88
import type { BillingEntity } from '@/lib/billing/core/usage-log'
@@ -32,6 +32,32 @@ function thresholdFor(percent: number): 0 | 80 | 100 {
3232
return 0
3333
}
3434

35+
/**
36+
* Replace the account's `limitNotifications` with `next` when `condition` holds, returning
37+
* whether the row was updated.
38+
*/
39+
async function writeLimitNotifications(
40+
scope: 'user' | 'organization',
41+
id: string,
42+
next: SQL,
43+
condition: SQL
44+
): Promise<boolean> {
45+
const written =
46+
scope === 'user'
47+
? await db
48+
.update(userStats)
49+
.set({ limitNotifications: next })
50+
.where(and(eq(userStats.userId, id), condition))
51+
.returning({ id: userStats.userId })
52+
: await db
53+
.update(organization)
54+
.set({ limitNotifications: next })
55+
.where(and(eq(organization.id, id), condition))
56+
.returning({ id: organization.id })
57+
58+
return written.length > 0
59+
}
60+
3561
/**
3662
* Atomically claim a threshold for a category: advance the stored value to
3763
* `threshold` only if it is currently lower, returning whether THIS call won the
@@ -50,20 +76,69 @@ async function claimThreshold(
5076
? sql`coalesce((${userStats.limitNotifications} ->> ${category})::int, 0) < ${threshold}`
5177
: sql`coalesce((${organization.limitNotifications} ->> ${category})::int, 0) < ${threshold}`
5278

53-
const claimed =
54-
scope === 'user'
79+
return writeLimitNotifications(scope, id, setExpr, onlyIfLower)
80+
}
81+
82+
const DAY_MS = 24 * 60 * 60 * 1000
83+
84+
/** One account's credits threshold, keyed on its billing period and limit. */
85+
export interface CreditsThresholdClaim {
86+
scope: 'user' | 'organization'
87+
id: string
88+
periodStart: Date
89+
limit: number
90+
threshold: 80 | 100
91+
}
92+
93+
/**
94+
* The stored claims for a credits threshold, and the condition under which it is still unclaimed:
95+
* `credits` holds the highest threshold emailed while `creditsPeriod` (the period's start day) and
96+
* `creditsLimit` (the limit in cents) still match, so a new period or a changed limit — in either
97+
* direction — re-arms both thresholds with no reset write, while within one a claim of 100 also
98+
* retires 80, and never the reverse.
99+
*/
100+
function creditsThresholdSql(claim: CreditsThresholdClaim) {
101+
const periodDay = Math.floor(claim.periodStart.getTime() / DAY_MS)
102+
const limitCents = Math.round(claim.limit * 100)
103+
const column =
104+
claim.scope === 'user' ? userStats.limitNotifications : organization.limitNotifications
105+
return {
106+
next: sql`coalesce(${column}, '{}'::jsonb) || jsonb_build_object('credits', ${claim.threshold}::int, 'creditsPeriod', ${periodDay}::bigint, 'creditsLimit', ${limitCents}::bigint)`,
107+
unclaimed: sql<boolean>`not (
108+
(${column} ->> 'creditsPeriod')::bigint is not distinct from ${periodDay}::bigint
109+
and (${column} ->> 'creditsLimit')::bigint is not distinct from ${limitCents}::bigint
110+
and coalesce((${column} ->> 'credits')::int, 0) >= ${claim.threshold}::int
111+
)`,
112+
}
113+
}
114+
115+
/**
116+
* Whether a credits threshold is still unclaimed, read with one indexed lookup of the account
117+
* row. A cheap pre-check only: callers run it on every completion above a threshold, so once
118+
* the threshold is claimed they stop there instead of resolving recipients. The atomic
119+
* {@link claimCreditsThreshold} remains the real dedup.
120+
*/
121+
export async function isCreditsThresholdUnclaimed(claim: CreditsThresholdClaim): Promise<boolean> {
122+
const { unclaimed } = creditsThresholdSql(claim)
123+
const [row] =
124+
claim.scope === 'user'
55125
? await db
56-
.update(userStats)
57-
.set({ limitNotifications: setExpr })
58-
.where(and(eq(userStats.userId, id), onlyIfLower))
59-
.returning({ id: userStats.userId })
126+
.select({ unclaimed })
127+
.from(userStats)
128+
.where(eq(userStats.userId, claim.id))
129+
.limit(1)
60130
: await db
61-
.update(organization)
62-
.set({ limitNotifications: setExpr })
63-
.where(and(eq(organization.id, id), onlyIfLower))
64-
.returning({ id: organization.id })
131+
.select({ unclaimed })
132+
.from(organization)
133+
.where(eq(organization.id, claim.id))
134+
.limit(1)
135+
return row?.unclaimed === true
136+
}
65137

66-
return claimed.length > 0
138+
/** Claim a credits threshold, returning whether THIS call won it. */
139+
export function claimCreditsThreshold(claim: CreditsThresholdClaim): Promise<boolean> {
140+
const { next, unclaimed } = creditsThresholdSql(claim)
141+
return writeLimitNotifications(claim.scope, claim.id, next, unclaimed)
67142
}
68143

69144
/** Re-arm a category (reset its stored threshold to 0) once usage falls back into the low band. */
@@ -108,7 +183,7 @@ async function isUnsubscribed(email: string): Promise<boolean> {
108183
* Returning an empty list means "nobody to notify" — the caller then skips the
109184
* claim so the dedup state isn't burned without an email going out.
110185
*/
111-
async function resolveRecipients(
186+
export async function resolveLimitEmailRecipients(
112187
scope: 'user' | 'organization',
113188
params: { userId?: string; userEmail?: string; userName?: string; organizationId?: string }
114189
): Promise<LimitEmailRecipient[]> {
@@ -207,7 +282,7 @@ export async function maybeSendLimitThresholdEmail(params: {
207282

208283
if (params.rearmOnly || desired === 0) return
209284

210-
const recipients = await resolveRecipients(scope, params)
285+
const recipients = await resolveLimitEmailRecipients(scope, params)
211286
if (recipients.length === 0) return
212287

213288
if (!(await claimThreshold(scope, stateId, category, desired))) return
Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/**
2+
* The shared reporting-usage read against a real ledger in a disposable PostgreSQL schema and a
3+
* real Redis. Skipped without `TEST_REDIS_URL`. Each test uses a fresh payer, so the in-process
4+
* cache is always cold and every read models a new process.
5+
*/
6+
7+
import { type AddressInfo, createServer, type Socket } from 'node:net'
8+
import type { db } from '@sim/db'
9+
import * as schema from '@sim/db/schema'
10+
import { readTestDatabaseUrl, readTestRedisUrl } from '@sim/db/testing/test-infrastructure'
11+
import { redisConfigMock, redisConfigMockFns } from '@sim/testing/mocks/redis-config.mock'
12+
import { generateId } from '@sim/utils/id'
13+
import { drizzle } from 'drizzle-orm/postgres-js'
14+
import Redis from 'ioredis'
15+
import postgres from 'postgres'
16+
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
17+
18+
const { transaction } = vi.hoisted(() => ({ transaction: vi.fn() }))
19+
const databaseUrl = readTestDatabaseUrl()
20+
const redisUrl = readTestRedisUrl()
21+
22+
vi.mock('@sim/db', () => ({ db: { transaction }, dbReplica: {} }))
23+
vi.mock('@/lib/core/config/redis', () => redisConfigMock)
24+
25+
import { readSoftGateUsageCost } from '@/lib/billing/core/reporting-usage-cache'
26+
import type { BillingEntity, UsageQueryPeriod } from '@/lib/billing/core/usage-log'
27+
28+
const schemaName = `reporting_usage_${generateId().replaceAll('-', '')}`
29+
const connection = postgres(databaseUrl, {
30+
max: 2,
31+
prepare: false,
32+
connection: { search_path: schemaName },
33+
onnotice: () => undefined,
34+
})
35+
const database = drizzle(connection, { schema }) as typeof db
36+
37+
const REPORTING: UsageQueryPeriod = {
38+
start: new Date('2026-01-01T00:00:00.000Z'),
39+
end: new Date('2027-01-01T00:00:00.000Z'),
40+
source: 'reporting',
41+
}
42+
43+
function sharedKey(payer: BillingEntity): string {
44+
return `usage:reporting:v1:${payer.type}:${payer.id}:reporting:${REPORTING.start.toISOString()}:${REPORTING.end.toISOString()}`
45+
}
46+
47+
describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
48+
let redis: Redis
49+
let payer: BillingEntity
50+
51+
beforeAll(async () => {
52+
redis = new Redis(redisUrl!, { lazyConnect: true, maxRetriesPerRequest: 0 })
53+
await redis.connect()
54+
await connection.unsafe(`CREATE SCHEMA "${schemaName}"`)
55+
await connection.unsafe(`CREATE TABLE usage_log (
56+
id text PRIMARY KEY, cost numeric NOT NULL, billing_entity_type text,
57+
billing_entity_id text, billing_period_start timestamp, billing_period_end timestamp,
58+
created_at timestamp NOT NULL
59+
)`)
60+
transaction.mockImplementation((callback) => database.transaction(callback))
61+
})
62+
63+
beforeEach(async () => {
64+
transaction.mockClear()
65+
redisConfigMockFns.mockGetRedisClient.mockReturnValue(redis)
66+
payer = { type: 'organization', id: generateId() }
67+
await connection`INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at)
68+
VALUES (${generateId()}, 4.25, 'organization', ${payer.id}, '2026-03-01'),
69+
(${generateId()}, 1.5, 'organization', ${payer.id}, '2026-06-01'),
70+
(${generateId()}, 99, 'organization', ${payer.id}, '2025-12-31')`
71+
})
72+
73+
afterEach(async () => {
74+
await redis.del(sharedKey(payer))
75+
})
76+
77+
afterAll(async () => {
78+
await redis?.quit()
79+
await connection.unsafe(`DROP SCHEMA IF EXISTS "${schemaName}" CASCADE`)
80+
await connection.end()
81+
})
82+
83+
it('serves a sum another process stored instead of the ledger', async () => {
84+
await redis.set(sharedKey(payer), '12.5', 'PX', 30_000)
85+
86+
await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(12.5)
87+
})
88+
89+
it('sums the ledger exactly on a miss and stores the sum for other processes', async () => {
90+
await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
91+
92+
await vi.waitFor(async () => expect(await redis.get(sharedKey(payer))).toBe('5.75'))
93+
const ttl = await redis.pttl(sharedKey(payer))
94+
expect(ttl).toBeGreaterThan(25_000)
95+
expect(ttl).toBeLessThanOrEqual(35_000)
96+
})
97+
98+
it('treats an unreadable stored sum as a miss', async () => {
99+
await redis.set(sharedKey(payer), 'not-a-number', 'PX', 30_000)
100+
101+
await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
102+
})
103+
104+
it('never lets a slower, older sum replace one stored while it ran', async () => {
105+
transaction.mockImplementationOnce(async (callback) => {
106+
const result = await database.transaction(callback)
107+
await redis.set(sharedKey(payer), '9.99', 'PX', 30_000)
108+
return result
109+
})
110+
111+
await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
112+
expect(await redis.get(sharedKey(payer))).toBe('9.99')
113+
})
114+
115+
it('sums the ledger when the Redis client cannot be built', async () => {
116+
redisConfigMockFns.mockGetRedisClient.mockImplementation(() => {
117+
throw new Error('Invalid Redis configuration')
118+
})
119+
120+
await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
121+
})
122+
123+
it('sums the ledger without issuing or queuing a command while Redis is not ready', async () => {
124+
const notReady = new Redis(redisUrl!, { lazyConnect: true, maxRetriesPerRequest: 0 })
125+
redisConfigMockFns.mockGetRedisClient.mockReturnValue(notReady)
126+
try {
127+
await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
128+
expect(notReady.status).toBe('wait')
129+
130+
await notReady.connect()
131+
await notReady.ping()
132+
expect(await redis.get(sharedKey(payer))).toBeNull()
133+
} finally {
134+
notReady.disconnect()
135+
}
136+
})
137+
138+
it('sums the ledger promptly when a connected Redis stops answering', async () => {
139+
const sockets = new Set<Socket>()
140+
/** Completes the client's handshake, then never answers a read. */
141+
const silent = createServer((socket) => {
142+
sockets.add(socket)
143+
socket.on('data', (data) => {
144+
const text = data.toString()
145+
if (/\bGET\b/i.test(text)) return
146+
socket.write('+OK\r\n'.repeat(text.match(/^\*\d+\r\n/gm)?.length ?? 0))
147+
})
148+
})
149+
await new Promise<void>((resolve) => silent.listen(0, '127.0.0.1', resolve))
150+
const { port } = silent.address() as AddressInfo
151+
const hung = new Redis({ host: '127.0.0.1', port, lazyConnect: true, enableReadyCheck: false })
152+
redisConfigMockFns.mockGetRedisClient.mockReturnValue(hung)
153+
try {
154+
await hung.connect()
155+
const startedAt = Date.now()
156+
await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
157+
expect(Date.now() - startedAt).toBeLessThan(2_000)
158+
} finally {
159+
hung.disconnect()
160+
for (const socket of sockets) socket.destroy()
161+
silent.close()
162+
}
163+
})
164+
})

0 commit comments

Comments
 (0)