From c12c0143ac1152239edf1b86beee11b24470724d Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 10 Sep 2026 02:14:12 -0700 Subject: [PATCH] fix(knowledge): cache the ingestion usage gate for a minute Every document run checked the payer's spend by summing the usage ledger for the billing period. The ledger grows with every indexed document, so a bulk connector sync re-read the whole period's ledger once per document: 41,000 sums over roughly 436,000 rows each during one crawl, which was the largest single database consumer while it ran and slowed unrelated queries, including search. Background ingestion now reads the gate through a per-process LRU cache keyed by payer, billing period and acting member, with a 60-second TTL and coalesced misses. Staleness fails in the harmless direction: a payer at their limit indexes for at most another minute, and a raised limit takes at most a minute to apply. Interactive callers keep reading the gate fresh. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01JBacX6HGVhPMUySuMfANwn --- .../billing/core/ingestion-usage-gate.test.ts | 98 +++++++++++++++++++ .../lib/billing/core/ingestion-usage-gate.ts | 69 +++++++++++++ .../documents/document-indexing-usage.test.ts | 2 + .../document-processing-source.test.ts | 2 + apps/sim/lib/knowledge/documents/service.ts | 4 +- 5 files changed, 173 insertions(+), 2 deletions(-) create mode 100644 apps/sim/lib/billing/core/ingestion-usage-gate.test.ts create mode 100644 apps/sim/lib/billing/core/ingestion-usage-gate.ts diff --git a/apps/sim/lib/billing/core/ingestion-usage-gate.test.ts b/apps/sim/lib/billing/core/ingestion-usage-gate.test.ts new file mode 100644 index 00000000000..dbecf11abf0 --- /dev/null +++ b/apps/sim/lib/billing/core/ingestion-usage-gate.test.ts @@ -0,0 +1,98 @@ +/** + * @vitest-environment node + */ +import { sleep } from '@sim/utils/helpers' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockCheck } = vi.hoisted(() => ({ mockCheck: vi.fn() })) + +vi.mock('@/lib/billing/core/billing-attribution', () => ({ + checkAttributedUsageLimits: mockCheck, +})) + +import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution' +import { + checkIngestionUsageLimits, + INGESTION_USAGE_GATE_TTL_MS, + resetIngestionUsageGateCache, +} from '@/lib/billing/core/ingestion-usage-gate' + +const ATTRIBUTION: BillingAttributionSnapshot = { + actorUserId: 'member-1', + workspaceId: null, + organizationId: 'org-1', + billedAccountUserId: 'owner-1', + billingEntity: { type: 'organization', id: 'org-1' }, + billingPeriod: { + start: '2026-09-01T00:00:00.000Z', + end: '2026-10-01T00:00:00.000Z', + source: 'stripe', + }, + payerSubscription: null, +} + +describe('checkIngestionUsageLimits', () => { + beforeEach(() => { + resetIngestionUsageGateCache() + mockCheck.mockReset().mockResolvedValue({ isExceeded: false }) + }) + afterEach(() => vi.restoreAllMocks()) + + it('reads the ledger once per payer, period and actor within the TTL', async () => { + await checkIngestionUsageLimits(ATTRIBUTION) + await checkIngestionUsageLimits(ATTRIBUTION) + await checkIngestionUsageLimits({ ...ATTRIBUTION, workspaceId: 'workspace-9' }) + expect(mockCheck).toHaveBeenCalledTimes(1) + }) + + it('collapses concurrent misses onto one ledger read', async () => { + let resolve!: (value: { isExceeded: boolean }) => void + mockCheck.mockReturnValueOnce(new Promise((r) => (resolve = r))) + const pending = Promise.all([ + checkIngestionUsageLimits(ATTRIBUTION), + checkIngestionUsageLimits(ATTRIBUTION), + checkIngestionUsageLimits(ATTRIBUTION), + ]) + await sleep(0) + resolve({ isExceeded: false }) + const results = await pending + expect(results.every((result) => result.isExceeded === false)).toBe(true) + expect(mockCheck).toHaveBeenCalledTimes(1) + }) + + it('keeps a refusal for the same bounded window, then reads fresh', async () => { + mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' }) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) + expect(mockCheck).toHaveBeenCalledTimes(1) + + /** `lru-cache` reads `performance.now()` and debounces it behind a real 1 ms timer. */ + const start = performance.now() + vi.spyOn(performance, 'now').mockReturnValue(start + INGESTION_USAGE_GATE_TTL_MS + 1) + await sleep(5) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false) + expect(mockCheck).toHaveBeenCalledTimes(2) + }) + + it('separates answers by actor, period and payer', async () => { + await checkIngestionUsageLimits(ATTRIBUTION) + await checkIngestionUsageLimits({ ...ATTRIBUTION, actorUserId: 'member-2' }) + await checkIngestionUsageLimits({ + ...ATTRIBUTION, + billingPeriod: { ...ATTRIBUTION.billingPeriod, start: '2026-10-01T00:00:00.000Z' }, + }) + await checkIngestionUsageLimits({ + ...ATTRIBUTION, + billedAccountUserId: 'owner-2', + billingEntity: { type: 'user', id: 'owner-2' }, + }) + expect(mockCheck).toHaveBeenCalledTimes(4) + }) + + it('does not cache a failed read', async () => { + mockCheck.mockRejectedValueOnce(new Error('ledger unavailable')) + await expect(checkIngestionUsageLimits(ATTRIBUTION)).rejects.toThrow('ledger unavailable') + await checkIngestionUsageLimits(ATTRIBUTION) + expect(mockCheck).toHaveBeenCalledTimes(2) + }) +}) diff --git a/apps/sim/lib/billing/core/ingestion-usage-gate.ts b/apps/sim/lib/billing/core/ingestion-usage-gate.ts new file mode 100644 index 00000000000..2f6f1c5f5f8 --- /dev/null +++ b/apps/sim/lib/billing/core/ingestion-usage-gate.ts @@ -0,0 +1,69 @@ +import { LRUCache } from 'lru-cache' +import { + type AttributedUsageLimitsResult, + type BillingAttributionSnapshot, + checkAttributedUsageLimits, +} from '@/lib/billing/core/billing-attribution' +import { coalesceLocally } from '@/lib/concurrency/singleflight' + +/** + * How long a usage-gate answer stays usable on the ingestion path. + * + * The gate sums the payer's usage ledger for the billing period, which grows + * with every indexed document, so a bulk sync that re-checks per document + * reads the whole period's ledger tens of thousands of times. Staleness fails + * in the harmless direction: a payer at their limit keeps indexing for at most + * this long, and a payer whose limit was just raised waits at most this long. + * Nothing on this path has a person waiting for the answer. + */ +export const INGESTION_USAGE_GATE_TTL_MS = 60 * 1000 + +/** Recent gate answers, with `LRUCache` supplying the TTL and the size bound. */ +const gateCache = new LRUCache({ + max: 10_000, + ttl: INGESTION_USAGE_GATE_TTL_MS, +}) + +/** + * The gate depends on who pays, for which period, and which member acts: the + * payer pool and the per-member cap are both part of the answer. + */ +function gateKey(attribution: BillingAttributionSnapshot): string { + return [ + attribution.billingEntity.type, + attribution.billingEntity.id, + attribution.billingPeriod.start, + attribution.billingPeriod.end, + attribution.billedAccountUserId, + attribution.actorUserId, + ].join(':') +} + +/** + * {@link checkAttributedUsageLimits} for background ingestion, with bounded + * staleness. Interactive callers (uploads, search, the settings surfaces) keep + * reading the gate fresh so a limit change is visible at once. + * + * `coalesceLocally` collapses the concurrent misses of a batch onto one ledger + * read and bounds a hung read at its settle deadline. The cache write stays on + * the value this caller received, so a producer that timed out and later + * resolved cannot overwrite a fresher answer. + */ +export async function checkIngestionUsageLimits( + attribution: BillingAttributionSnapshot +): Promise { + const key = gateKey(attribution) + const cached = gateCache.get(key) + if (cached !== undefined) return cached + + const result = await coalesceLocally(`ingestion-usage-gate:${key}`, () => + checkAttributedUsageLimits(attribution) + ) + gateCache.set(key, result) + return result +} + +/** Drops every cached gate answer. Test seam; never called in production code. */ +export function resetIngestionUsageGateCache(): void { + gateCache.clear() +} diff --git a/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts b/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts index e68b00db64f..437de7bf205 100644 --- a/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts +++ b/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts @@ -65,11 +65,13 @@ vi.mock('@/providers/utils', () => ({ })) import * as billingAttribution from '@/lib/billing/core/billing-attribution' +import { resetIngestionUsageGateCache } from '@/lib/billing/core/ingestion-usage-gate' import * as embeddingClient from '@/lib/embeddings/client' import { processDocumentAsync } from '@/lib/knowledge/documents/service' const mockEmbeddingCapacity = vi.fn() beforeEach(() => { + resetIngestionUsageGateCache() vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation( mockCheckAttributedUsageLimits ) diff --git a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts index c3eeff755ca..5b36cbdf3c2 100644 --- a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts +++ b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts @@ -66,6 +66,7 @@ vi.mock('@/lib/uploads/server/metadata', () => ({ })) import * as billingAttribution from '@/lib/billing/core/billing-attribution' +import { resetIngestionUsageGateCache } from '@/lib/billing/core/ingestion-usage-gate' import { env } from '@/lib/core/config/env' import { markInsideTriggerRun, @@ -92,6 +93,7 @@ import { MAX_PROCESSING_ATTEMPTS } from '@/lib/knowledge/documents/types' const mockEmbeddingCapacity = vi.fn() beforeEach(() => { + resetIngestionUsageGateCache() vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation( mockCheckAttributedUsageLimits ) diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 8a555fb1834..7b962342a17 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -34,9 +34,9 @@ import { assertBillingAttributionOwner, assertBillingAttributionSnapshot, type BillingAttributionSnapshot, - checkAttributedUsageLimits, toBillingContext, } from '@/lib/billing/core/billing-attribution' +import { checkIngestionUsageLimits } from '@/lib/billing/core/ingestion-usage-gate' import { recordUsage } from '@/lib/billing/core/usage-log' import { applyStorageUsageDeltasInTx, @@ -1618,7 +1618,7 @@ export async function processDocumentAsync( assertBillingAttributionOwner(billingAttribution, ctx) const documentActorUserId = billingAttribution.actorUserId - const usageGate = await checkAttributedUsageLimits(billingAttribution) + const usageGate = await checkIngestionUsageLimits(billingAttribution) if (usageGate.isExceeded) { logger.warn(`[${documentId}] Usage limit reached — skipping document indexing`) throw new UsageLimitDocumentProcessingError(