From fda1111ca263cb03b94c458b498db7dd15f919bf Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 12 Oct 2026 14:23:38 -0700 Subject: [PATCH 1/5] improvement(rate-limit): admit provider requests in one statement on the database bucket Lock, refill, decide and spend every dimension in a single CTE so a hot bucket row stays locked for one statement plus the commit instead of a round trip per dimension. The clock is read before the statement, matching the Redis backend, so a contended admission never sees more tokens. --- .../rate-limiter/provider-admission.test.ts | 16 +- .../storage/db-token-bucket.test.ts | 35 --- .../rate-limiter/storage/db-token-bucket.ts | 258 +++++++++++++----- .../provider-admission.integration.ts | 31 +++ 4 files changed, 234 insertions(+), 106 deletions(-) diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts index 180306504c9..8bd52a15c51 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts @@ -167,8 +167,8 @@ describe('provider admission', () => { maxWaitMs: 1, } let stored: { key: string; tokens: string; lastRefillAt: Date } | undefined - dbChainMockFns.values.mockImplementation((rows) => { - stored ??= rows.find((row: { key: string }) => row.key.endsWith(':requests')) + dbChainMockFns.values.mockImplementation((row) => { + stored ??= { ...row } return { onConflictDoNothing: vi.fn().mockResolvedValue(undefined) } }) dbChainMockFns.limit.mockImplementation(async () => [stored]) @@ -177,9 +177,15 @@ describe('provider admission', () => { return { where: vi.fn().mockResolvedValue(undefined) } }) const bucket = new DbTokenBucket() - consumeTokens.mockImplementation((reservations, options) => - bucket.consumeTokensAtomically(reservations, options) - ) + /** The single-dimension refill arithmetic the atomic statement mirrors in SQL. */ + consumeTokens.mockImplementation(async ([reservation]) => { + const result = await bucket.consumeTokens( + reservation.key, + reservation.cost, + reservation.config + ) + return { allowed: result.allowed, retryAfterMs: result.retryAfterMs ?? 0 } + }) for (let request = 0; request < 16; request++) await waitForProviderAdmission(input) await expect(waitForProviderAdmission(input)).rejects.toMatchObject({ retryAfterMs: 1000 }) diff --git a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts index 6c1655d7112..443f3dc8100 100644 --- a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts +++ b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts @@ -44,41 +44,6 @@ describe('PostgreSQL token bucket', () => { expect(dbChainMockFns.set).toHaveBeenCalledWith(expect.objectContaining({ tokens: '0' })) }) - it('initializes request and cooldown buckets in a consistent order without duplicate keys', async () => { - const now = new Date() - dbChainMockFns.limit.mockResolvedValue([ - { key: 'cooldown', tokens: '0', lastRefillAt: now, blockedUntil: null }, - { key: 'requests', tokens: '10', lastRefillAt: now, blockedUntil: null }, - { key: 'tokens', tokens: '10', lastRefillAt: now, blockedUntil: null }, - ]) - - expect( - await new DbTokenBucket().consumeTokensAtomically( - [ - { key: 'tokens', cost: 2, config: CONFIG }, - { key: 'requests', cost: 1, config: CONFIG }, - ], - { cooldownKeys: ['cooldown', 'cooldown'], deadlineAt: now.getTime() + 1000 } - ) - ).toEqual({ allowed: true, retryAfterMs: 0 }) - - expect(dbChainMockFns.values).toHaveBeenCalledExactlyOnceWith([ - { key: 'cooldown', tokens: '0', lastRefillAt: now, updatedAt: now }, - { key: 'requests', tokens: '10', lastRefillAt: now, updatedAt: now }, - { key: 'tokens', tokens: '10', lastRefillAt: now, updatedAt: now }, - ]) - expect(dbChainMockFns.set).toHaveBeenCalledWith({ - tokens: '8', - lastRefillAt: now, - updatedAt: now, - }) - expect(dbChainMockFns.set).toHaveBeenCalledWith({ - tokens: '9', - lastRefillAt: now, - updatedAt: now, - }) - }) - it('accepts an empty reservation without attempting an empty insert', async () => { dbChainMockFns.limit.mockResolvedValue([]) diff --git a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts index 6f818980916..81453583d5f 100644 --- a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts +++ b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts @@ -1,6 +1,6 @@ import { db } from '@sim/db' import { rateLimitBucket } from '@sim/db/schema' -import { eq, inArray, sql } from 'drizzle-orm' +import { eq, type SQL, sql } from 'drizzle-orm' import type { AtomicAdmissionOptions, AtomicAdmissionResult, @@ -11,6 +11,166 @@ import type { TokenStatus, } from '@/lib/core/rate-limiter/storage/adapter' +/** Thrown inside the admission transaction to roll back a spend that finished past its deadline. */ +class AdmissionExpiredError extends Error {} + +interface AdmissionRow extends Record { + kind: 'reservation' | 'cooldown' + cost: number | null + available: number | null + last_refill_ms: number | null + refill_rate: number | null + refill_interval_ms: number | null + blocked_until_ms: number | null + allowed: boolean + spent: number +} + +/** An array of scalar binds; the pool cannot bind a JavaScript array as one parameter. */ +function sqlArray(values: readonly (string | number)[], type: 'text' | 'float8'): SQL { + return sql`ARRAY[${sql.join( + values.map((value) => sql`${value}`), + sql`, ` + )}]::${sql.raw(type)}[]` +} + +/** + * Locks every key in sorted order, refills each reservation to `now` with the same float8 + * arithmetic {@link DbTokenBucket.consumeTokens} uses in JavaScript, and spends all reservations + * only when no cooldown is active and every balance covers its cost. Returns one row per + * reservation and per cooldown key, each carrying the decision and how many buckets were spent. + */ +function buildAdmissionStatement( + reservations: readonly TokenBucketReservation[], + cooldownKeys: readonly string[], + keys: readonly string[], + now: number +): SQL { + return sql` + WITH reservation AS ( + SELECT * FROM unnest( + ${sqlArray( + reservations.map((item) => item.key), + 'text' + )}, + ${sqlArray( + reservations.map((item) => item.cost), + 'float8' + )}, + ${sqlArray( + reservations.map((item) => item.config.maxTokens), + 'float8' + )}, + ${sqlArray( + reservations.map((item) => item.config.refillRate), + 'float8' + )}, + ${sqlArray( + reservations.map((item) => item.config.refillIntervalMs), + 'float8' + )} + ) AS r(key, cost, max_tokens, refill_rate, refill_interval_ms) + ), + locked AS ( + SELECT key, tokens, last_refill_at, blocked_until + FROM ${rateLimitBucket} + WHERE key = ANY(${sqlArray(keys, 'text')}) + ORDER BY key + FOR UPDATE + ), + balance AS ( + SELECT + r.key, + r.cost, + r.refill_rate, + r.refill_interval_ms, + LEAST(r.max_tokens, GREATEST(0, l.tokens::float8) + e.elapsed * r.refill_rate) AS available, + m.last_ms + e.elapsed * r.refill_interval_ms AS last_refill_ms, + l.last_refill_at + (e.elapsed * r.refill_interval_ms) * interval '1 millisecond' AS last_refill_at + FROM reservation r + JOIN locked l ON l.key = r.key + CROSS JOIN LATERAL ( + SELECT (extract(epoch FROM l.last_refill_at) * 1000)::float8 AS last_ms + ) m + CROSS JOIN LATERAL ( + SELECT GREATEST(0, floor((${now}::float8 - m.last_ms) / r.refill_interval_ms)) AS elapsed + ) e + ), + cooldown AS ( + SELECT (extract(epoch FROM blocked_until) * 1000)::float8 AS blocked_until_ms + FROM locked + WHERE key = ANY(${sqlArray(cooldownKeys, 'text')}) + ), + decision AS ( + SELECT + NOT EXISTS (SELECT 1 FROM balance WHERE available < cost) + AND NOT EXISTS (SELECT 1 FROM cooldown WHERE blocked_until_ms > ${now}::float8) AS allowed + ), + spent AS ( + UPDATE ${rateLimitBucket} AS bucket + SET + tokens = (balance.available - balance.cost)::text::numeric, + last_refill_at = balance.last_refill_at, + updated_at = timestamp '1970-01-01' + ${now}::float8 * interval '1 millisecond' + FROM balance, decision + WHERE decision.allowed AND bucket.key = balance.key + RETURNING bucket.key + ), + outcome AS ( + SELECT decision.allowed, (SELECT count(*) FROM spent)::int AS spent FROM decision + ) + SELECT 'reservation' AS kind, b.cost, b.available, b.last_refill_ms, b.refill_rate, + b.refill_interval_ms, NULL::float8 AS blocked_until_ms, o.allowed, o.spent + FROM balance b CROSS JOIN outcome o + UNION ALL + SELECT 'cooldown', NULL, NULL, NULL, NULL, NULL, c.blocked_until_ms, o.allowed, o.spent + FROM cooldown c CROSS JOIN outcome o + ` +} + +/** + * Derives the wait from the locked balances and cross-checks it against the statement's own + * decision, so the JavaScript reading and the SQL spend can never disagree silently: a mismatch, + * or a reservation whose bucket vanished, throws and rolls the spend back. + */ +function readAdmissionResult( + rows: readonly AdmissionRow[], + reservations: readonly TokenBucketReservation[], + now: number +): AtomicAdmissionResult { + let retryAfterMs = 0 + let reserved = 0 + for (const row of rows) { + if (row.kind === 'cooldown') { + if (row.blocked_until_ms !== null) { + retryAfterMs = Math.max(retryAfterMs, Number(row.blocked_until_ms) - now) + } + continue + } + reserved++ + const available = Number(row.available) + const cost = Number(row.cost) + if (available < cost) { + retryAfterMs = Math.max( + retryAfterMs, + Number(row.last_refill_ms) + + Math.ceil((cost - available) / Number(row.refill_rate)) * Number(row.refill_interval_ms) - + now + ) + } + } + if (reserved !== reservations.length) { + throw new Error('Rate limit bucket disappeared during admission') + } + const allowed = rows[0]?.allowed ?? true + const spent = rows[0]?.spent ?? 0 + const spentExpected = allowed ? new Set(reservations.map((item) => item.key)).size : 0 + if (allowed !== retryAfterMs <= 0 || spent !== spentExpected) { + throw new Error('Rate limit admission decision is inconsistent') + } + return allowed ? { allowed: true, retryAfterMs: 0 } : { allowed: false, retryAfterMs } +} + export class DbTokenBucket implements RateLimitStorageAdapter { async consumeTokens( key: string, @@ -78,23 +238,32 @@ export class DbTokenBucket implements RateLimitStorageAdapter { }) } - /** Locks all dimensions and provider gates before spending any capacity. */ + /** + * Locks all dimensions and provider gates before spending any capacity. + * + * The lock, refill, decision, and spend run as one statement, so the bucket rows stay locked + * for that statement and the commit rather than for a round trip per dimension: every waiter + * on a hot provider key queues behind that whole hold. The clock is read just before the + * statement, so time spent waiting for the lock is not refilled; a contended admission can + * only see fewer tokens than it would have read after the wait, never more. + */ async consumeTokensAtomically( reservations: readonly TokenBucketReservation[], options: AtomicAdmissionOptions ): Promise { options.signal?.throwIfAborted() if (Date.now() >= options.deadlineAt) return { allowed: false, retryAfterMs: 0 } - return db.transaction(async (tx) => { - const remainingMs = Math.max(1, options.deadlineAt - Date.now()) - await tx.execute( - sql`SELECT set_config('statement_timeout', ${String(remainingMs)}, true), set_config('lock_timeout', ${String(remainingMs)}, true)` - ) - const keys = [ - ...new Set([...options.cooldownKeys, ...reservations.map((item) => item.key)]), - ].sort() - const createdAt = new Date() - if (keys.length > 0) { + const keys = [ + ...new Set([...options.cooldownKeys, ...reservations.map((item) => item.key)]), + ].sort() + if (keys.length === 0) return { allowed: true, retryAfterMs: 0 } + try { + return await db.transaction(async (tx) => { + const remainingMs = Math.max(1, options.deadlineAt - Date.now()) + await tx.execute( + sql`SELECT set_config('statement_timeout', ${String(remainingMs)}, true), set_config('lock_timeout', ${String(remainingMs)}, true)` + ) + const createdAt = new Date() await tx .insert(rateLimitBucket) .values( @@ -109,62 +278,19 @@ export class DbTokenBucket implements RateLimitStorageAdapter { }) ) .onConflictDoNothing() - } - const rows = await tx - .select() - .from(rateLimitBucket) - .where(inArray(rateLimitBucket.key, keys)) - .orderBy(rateLimitBucket.key) - .for('update') - .limit(keys.length) - options.signal?.throwIfAborted() - const now = Date.now() - if (now >= options.deadlineAt) return { allowed: false, retryAfterMs: 0 } - let retryAfterMs = 0 - for (const key of options.cooldownKeys) { - const row = rows.find((candidate) => candidate.key === key) - retryAfterMs = Math.max(retryAfterMs, (row?.blockedUntil?.getTime() ?? 0) - now) - } - const balances = reservations.map((reservation) => { - const row = rows.find((candidate) => candidate.key === reservation.key) - if (!row) throw new Error('Rate limit bucket disappeared during admission') - const { config, cost } = reservation - const elapsed = Math.max( - 0, - Math.floor((now - row.lastRefillAt.getTime()) / config.refillIntervalMs) - ) - const available = Math.min( - config.maxTokens, - Math.max(0, Number(row.tokens)) + elapsed * config.refillRate + options.signal?.throwIfAborted() + const now = Date.now() + const rows = await tx.execute( + buildAdmissionStatement(reservations, options.cooldownKeys, keys, now) ) - const lastRefillAt = new Date( - row.lastRefillAt.getTime() + elapsed * config.refillIntervalMs - ) - if (available < cost) { - retryAfterMs = Math.max( - retryAfterMs, - lastRefillAt.getTime() + - Math.ceil((cost - available) / config.refillRate) * config.refillIntervalMs - - now - ) - } - return { reservation, available, lastRefillAt } + options.signal?.throwIfAborted() + if (Date.now() >= options.deadlineAt) throw new AdmissionExpiredError() + return readAdmissionResult(rows, reservations, now) }) - if (retryAfterMs > 0) return { allowed: false, retryAfterMs } - options.signal?.throwIfAborted() - for (const { reservation, available, lastRefillAt } of balances) { - await tx - .update(rateLimitBucket) - .set({ - tokens: String(available - reservation.cost), - lastRefillAt, - updatedAt: new Date(now), - }) - .where(eq(rateLimitBucket.key, reservation.key)) - } - options.signal?.throwIfAborted() - return { allowed: true, retryAfterMs: 0 } - }) + } catch (error) { + if (error instanceof AdmissionExpiredError) return { allowed: false, retryAfterMs: 0 } + throw error + } } async getCooldownUntil(key: string): Promise { diff --git a/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts b/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts index b2fa4386a74..439d5450a54 100644 --- a/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts @@ -96,6 +96,37 @@ describe.each(['PostgreSQL', 'Redis'] as const)('%s provider admission', (backen }) }) + it('refills fractional reservations by whole intervals and returns the full wait', async () => { + const reservation = { + key: `${key}:requests`, + cost: 1, + config: { maxTokens: 1, refillRate: 0.25, refillIntervalMs: 1000 }, + } + const options = { cooldownKeys: [`${key}:cooldown`], deadlineAt: now.getTime() + 600_000 } + expect(await first.consumeTokensAtomically([reservation], options)).toEqual({ + allowed: true, + retryAfterMs: 0, + }) + vi.setSystemTime(now.getTime() + 1500) + expect(await second.consumeTokensAtomically([reservation], options)).toEqual({ + allowed: false, + retryAfterMs: 2500, + }) + expect(await first.getTokenStatus(reservation.key, reservation.config)).toMatchObject({ + tokensAvailable: 0.25, + lastRefillAt: new Date(now.getTime() + 1000), + }) + vi.setSystemTime(now.getTime() + 4000) + expect(await second.consumeTokensAtomically([reservation], options)).toEqual({ + allowed: true, + retryAfterMs: 0, + }) + expect(await first.getTokenStatus(reservation.key, reservation.config)).toMatchObject({ + tokensAvailable: 0, + lastRefillAt: new Date(now.getTime() + 4000), + }) + }) + it.each([ { state: 'new', spentRequests: 0, spentTokens: 0, admitted: 8, remainingTokens: 20 }, { state: 'existing', spentRequests: 2, spentTokens: 20, admitted: 6, remainingTokens: 20 }, From 4542b2a6cf1030cd111d492ada1c35f88f8c762c Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 12 Oct 2026 14:23:39 -0700 Subject: [PATCH 2/5] improvement(connectors): skip the live document count on unfinished sync slices An unfinished content pass keeps the count its last finished pass recorded, the pass lastSyncAt still names. Connectors with no recorded count are still counted so a first pass shows progress. --- .../knowledge/connectors/sync-engine.test.ts | 44 +++++++++++++++++++ .../lib/knowledge/connectors/sync-engine.ts | 34 +++++++++----- 2 files changed, 67 insertions(+), 11 deletions(-) diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index a9e18722987..29fe7aaf43b 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -1591,6 +1591,50 @@ describe('completeSuccessfulSync', () => { expect(successWrite).toMatchObject({ consecutiveFailures: 0 }) }) + it.each([ + { recordedDocCount: 7, written: undefined }, + { recordedDocCount: null, written: 4 }, + ])( + 'leaves an unfinished pass on the last finished count unless none was recorded: $recordedDocCount', + async ({ recordedDocCount, written }) => { + const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine') + queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1', lastSyncDocCount: 7 }]) + queueTableRows(schemaMock.document, [{ count: 4 }]) + dbChainMockFns.returning + .mockResolvedValueOnce([{ id: 'log-1' }]) + .mockResolvedValueOnce([{ id: 'c-1' }]) + + await expect( + completeSuccessfulSync( + 'c-1', + 'kb-1', + 'log-1', + 60, + { ...RESULT, docsFailed: 0 }, + null, + { + complete: false, + checkpoint: { + unsafe: false, + startedAt: '2026-09-04T00:00:00Z', + listedCount: 25, + resumeAt: null, + }, + }, + null, + false, + recordedDocCount + ) + ).resolves.toBe(true) + + const successWrite = dbChainMockFns.set.mock.calls + .map(([value]) => value as Record) + .find((value) => value.status === 'active') + expect(successWrite?.lastSyncDocCount).toBe(written) + } + ) + it('publishes neither terminal state when lock ownership is gone', async () => { const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine') diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.ts b/apps/sim/lib/knowledge/connectors/sync-engine.ts index 10f467de4a1..34202895fe7 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.ts @@ -399,7 +399,9 @@ export async function completeSuccessfulSync( contentPass?: ContentPassOutcome, directoryNotice: string | null = null, /** A pending ACL rewrite this run began but did not finish: the flag stays and the next run resumes it. */ - accessRewriteUnfinished = false + accessRewriteUnfinished = false, + /** The count the connector held when this run took its lock. */ + recordedDocCount: number | null = null ): Promise { const processingDispatchFailed = result.processingDispatch.failed > 0 const contentNotice = @@ -431,16 +433,25 @@ export async function completeSuccessfulSync( * owned count the next pass's mass-deletion guard starts from, so a tally of this run's * adds and deletes would drift; when it cannot be read the previous count stands, which the * guard already treats as a floor. + * + * A slice that leaves its content pass unfinished keeps the count the last finished pass + * recorded, the same pass `lastSyncAt` still names, so a large connector is not counted again + * on every slice of a pass whose count would only be replaced when it finishes. A connector + * with no count yet is still counted, so a first pass shows its progress. */ - const actualDocCount = await countLiveConnectorDocuments(connectorId).catch((error: unknown) => { - const diagnostic = getConnectorFailureDiagnostic(error) - if (diagnostic?.category !== 'database') throw error - logger.warn('Could not count connector documents; keeping the previous count', { - connectorId, - diagnostic, - }) - return null - }) + const keepRecordedCount = + contentPass !== undefined && !contentPass.complete && recordedDocCount !== null + const actualDocCount = keepRecordedCount + ? null + : await countLiveConnectorDocuments(connectorId).catch((error: unknown) => { + const diagnostic = getConnectorFailureDiagnostic(error) + if (diagnostic?.category !== 'database') throw error + logger.warn('Could not count connector documents; keeping the previous count', { + connectorId, + diagnostic, + }) + return null + }) try { return await db.transaction(async (tx) => { const [lockedKnowledgeBase] = await tx @@ -1417,7 +1428,8 @@ export async function executeSync( reconciliationHoldNotice, contentPass, directoryNotice, - accessRewriteUnfinished + accessRewriteUnfinished, + connector.lastSyncDocCount ) if (!completionLanded) { From 67ec3a8100f7dfa93b3f531dd38b6c745c2dadfd Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 12 Oct 2026 14:43:37 -0700 Subject: [PATCH 3/5] fix(rate-limit): judge atomic admission at the moment the bucket locks are held Advance the caller's clock by the lock wait measured inside the statement, so a contended admission refills the time it spent waiting as it did before the single-statement change. Cover the spend rollback when the deadline passes or the caller aborts while the statement runs. --- .../rate-limiter/storage/db-token-bucket.ts | 54 ++++--- .../provider-admission.integration.ts | 142 +++++++++++++++++- 2 files changed, 177 insertions(+), 19 deletions(-) diff --git a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts index 81453583d5f..d7adbd674e7 100644 --- a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts +++ b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts @@ -24,6 +24,7 @@ interface AdmissionRow extends Record { blocked_until_ms: number | null allowed: boolean spent: number + now_ms: number } /** An array of scalar binds; the pool cannot bind a JavaScript array as one parameter. */ @@ -35,16 +36,23 @@ function sqlArray(values: readonly (string | number)[], type: 'text' | 'float8') } /** - * Locks every key in sorted order, refills each reservation to `now` with the same float8 - * arithmetic {@link DbTokenBucket.consumeTokens} uses in JavaScript, and spends all reservations - * only when no cooldown is active and every balance covers its cost. Returns one row per - * reservation and per cooldown key, each carrying the decision and how many buckets were spent. + * Locks every key in sorted order, refills each reservation to the moment every lock is held + * with the same float8 arithmetic {@link DbTokenBucket.consumeTokens} uses in JavaScript, and + * spends all reservations only when no cooldown is active and every balance covers its cost. + * Returns one row per reservation and per cooldown key, each carrying the decision, how many + * buckets were spent, and the clock the decision used. + * + * That clock is the caller's `sentAt` advanced by the time the statement spent acquiring its + * locks, in whole milliseconds: `clock` counts every `locked` row before reading + * `clock_timestamp()`, so the read happens only once the last lock is held. Measuring the wait + * on the database rather than reading its absolute time keeps every stored instant on the + * application clock that wrote it. */ function buildAdmissionStatement( reservations: readonly TokenBucketReservation[], cooldownKeys: readonly string[], keys: readonly string[], - now: number + sentAt: number ): SQL { return sql` WITH reservation AS ( @@ -78,6 +86,11 @@ function buildAdmissionStatement( ORDER BY key FOR UPDATE ), + clock AS ( + SELECT ${sentAt}::float8 + + floor(extract(epoch FROM (held.locked_at - statement_timestamp())) * 1000) AS now_ms + FROM (SELECT clock_timestamp() AS locked_at FROM (SELECT count(*) FROM locked) every_lock) held + ), balance AS ( SELECT r.key, @@ -89,11 +102,12 @@ function buildAdmissionStatement( l.last_refill_at + (e.elapsed * r.refill_interval_ms) * interval '1 millisecond' AS last_refill_at FROM reservation r JOIN locked l ON l.key = r.key + CROSS JOIN clock CROSS JOIN LATERAL ( SELECT (extract(epoch FROM l.last_refill_at) * 1000)::float8 AS last_ms ) m CROSS JOIN LATERAL ( - SELECT GREATEST(0, floor((${now}::float8 - m.last_ms) / r.refill_interval_ms)) AS elapsed + SELECT GREATEST(0, floor((clock.now_ms - m.last_ms) / r.refill_interval_ms)) AS elapsed ) e ), cooldown AS ( @@ -104,26 +118,30 @@ function buildAdmissionStatement( decision AS ( SELECT NOT EXISTS (SELECT 1 FROM balance WHERE available < cost) - AND NOT EXISTS (SELECT 1 FROM cooldown WHERE blocked_until_ms > ${now}::float8) AS allowed + AND NOT EXISTS ( + SELECT 1 FROM cooldown, clock WHERE cooldown.blocked_until_ms > clock.now_ms + ) AS allowed ), spent AS ( UPDATE ${rateLimitBucket} AS bucket SET tokens = (balance.available - balance.cost)::text::numeric, last_refill_at = balance.last_refill_at, - updated_at = timestamp '1970-01-01' + ${now}::float8 * interval '1 millisecond' - FROM balance, decision + updated_at = timestamp '1970-01-01' + clock.now_ms * interval '1 millisecond' + FROM balance, decision, clock WHERE decision.allowed AND bucket.key = balance.key RETURNING bucket.key ), outcome AS ( - SELECT decision.allowed, (SELECT count(*) FROM spent)::int AS spent FROM decision + SELECT decision.allowed, (SELECT count(*) FROM spent)::int AS spent, clock.now_ms + FROM decision, clock ) SELECT 'reservation' AS kind, b.cost, b.available, b.last_refill_ms, b.refill_rate, - b.refill_interval_ms, NULL::float8 AS blocked_until_ms, o.allowed, o.spent + b.refill_interval_ms, NULL::float8 AS blocked_until_ms, o.allowed, o.spent, o.now_ms FROM balance b CROSS JOIN outcome o UNION ALL - SELECT 'cooldown', NULL, NULL, NULL, NULL, NULL, c.blocked_until_ms, o.allowed, o.spent + SELECT 'cooldown', NULL, NULL, NULL, NULL, NULL, c.blocked_until_ms, o.allowed, o.spent, + o.now_ms FROM cooldown c CROSS JOIN outcome o ` } @@ -243,9 +261,8 @@ export class DbTokenBucket implements RateLimitStorageAdapter { * * The lock, refill, decision, and spend run as one statement, so the bucket rows stay locked * for that statement and the commit rather than for a round trip per dimension: every waiter - * on a hot provider key queues behind that whole hold. The clock is read just before the - * statement, so time spent waiting for the lock is not refilled; a contended admission can - * only see fewer tokens than it would have read after the wait, never more. + * on a hot provider key queues behind that whole hold. Refill and cooldowns are judged at the + * moment the locks are held, as before, so a wait for a contended lock is refilled. */ async consumeTokensAtomically( reservations: readonly TokenBucketReservation[], @@ -279,12 +296,13 @@ export class DbTokenBucket implements RateLimitStorageAdapter { ) .onConflictDoNothing() options.signal?.throwIfAborted() - const now = Date.now() + const sentAt = Date.now() const rows = await tx.execute( - buildAdmissionStatement(reservations, options.cooldownKeys, keys, now) + buildAdmissionStatement(reservations, options.cooldownKeys, keys, sentAt) ) + const now = rows[0]?.now_ms ?? sentAt options.signal?.throwIfAborted() - if (Date.now() >= options.deadlineAt) throw new AdmissionExpiredError() + if (Math.max(now, Date.now()) >= options.deadlineAt) throw new AdmissionExpiredError() return readAdmissionResult(rows, reservations, now) }) } catch (error) { diff --git a/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts b/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts index 439d5450a54..421650dea1c 100644 --- a/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts @@ -1,8 +1,15 @@ +import { db } from '@sim/db' +import { rateLimitBucket } from '@sim/db/schema' import { readTestRedisUrl } from '@sim/db/testing/test-infrastructure' +import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' +import { eq } from 'drizzle-orm' import Redis from 'ioredis' import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' -import type { RateLimitStorageAdapter } from '@/lib/core/rate-limiter/storage/adapter' +import type { + AtomicAdmissionResult, + RateLimitStorageAdapter, +} from '@/lib/core/rate-limiter/storage/adapter' import { DbTokenBucket } from '@/lib/core/rate-limiter/storage/db-token-bucket' import { RedisTokenBucket } from '@/lib/core/rate-limiter/storage/redis-token-bucket' @@ -249,3 +256,136 @@ describe.each(['PostgreSQL', 'Redis'] as const)('%s provider admission', (backen }) }) }) + +describe('PostgreSQL admission under a held bucket lock', () => { + const bucket = new DbTokenBucket() + let key: string + let now: number + + /** How long the competing transaction holds the row; long enough to cross a refill interval. */ + const LOCK_HOLD_MS = 1200 + + beforeEach(() => { + now = Date.now() + vi.useFakeTimers({ toFake: ['Date'] }) + vi.setSystemTime(now) + key = `provider-integration:${generateId()}` + }) + + afterEach(async () => { + vi.useRealTimers() + await bucket.resetBucket(key) + }) + + /** Locks the bucket row from another connection, runs `whileHeld`, then commits the holder. */ + async function withRowLockHeld(whileHeld: () => Promise): Promise { + let held: () => void = () => {} + let release: () => void = () => {} + const acquired = new Promise((resolve) => { + held = resolve + }) + const released = new Promise((resolve) => { + release = resolve + }) + const holder = db.transaction(async (tx) => { + await tx.select().from(rateLimitBucket).where(eq(rateLimitBucket.key, key)).for('update') + held() + await released + }) + await acquired + try { + await whileHeld() + } finally { + release() + await holder + } + } + + /** + * Starts an admission and resolves once it is queued behind the held row lock. The result is + * wrapped so awaiting this does not also await the blocked admission. + */ + async function queuedBehindLock( + admit: () => Promise + ): Promise<{ result: Promise }> { + const result = admit() + result.catch(() => {}) + await sleep(300) + return { result } + } + + async function storedTokens(): Promise { + const [row] = await db + .select({ tokens: rateLimitBucket.tokens }) + .from(rateLimitBucket) + .where(eq(rateLimitBucket.key, key)) + return Number(row?.tokens) + } + + it('refills the time spent waiting for the lock and measures the wait from then', async () => { + const reservation = { + key, + cost: 2, + config: { maxTokens: 2, refillRate: 1, refillIntervalMs: 1000 }, + } + const options = { cooldownKeys: [], deadlineAt: now + 60_000 } + expect(await bucket.consumeTokensAtomically([reservation], options)).toEqual({ + allowed: true, + retryAfterMs: 0, + }) + let denied: { result: Promise } | undefined + await withRowLockHeld(async () => { + denied = await queuedBehindLock(() => bucket.consumeTokensAtomically([reservation], options)) + await sleep(LOCK_HOLD_MS - 300) + }) + const result = await denied?.result + expect(result?.allowed).toBe(false) + expect(result?.retryAfterMs).toBeGreaterThan(0) + expect(result?.retryAfterMs).toBeLessThan(1000) + + let admitted: { result: Promise } | undefined + await withRowLockHeld(async () => { + admitted = await queuedBehindLock(() => + bucket.consumeTokensAtomically([{ ...reservation, cost: 1 }], options) + ) + await sleep(LOCK_HOLD_MS - 300) + }) + expect(await admitted?.result).toEqual({ allowed: true, retryAfterMs: 0 }) + }) + + it.each(['deadline', 'abort'] as const)( + 'rolls back a spend whose %s arrives while the statement runs', + async (interruption) => { + const reservation = { + key, + cost: 5, + config: { maxTokens: 10, refillRate: 1, refillIntervalMs: 600_000 }, + } + expect( + await bucket.consumeTokensAtomically([{ ...reservation, cost: 0 }], { + cooldownKeys: [], + deadlineAt: now + 60_000, + }) + ).toEqual({ allowed: true, retryAfterMs: 0 }) + const controller = new AbortController() + let pending: { result: Promise } | undefined + await withRowLockHeld(async () => { + pending = await queuedBehindLock(() => + bucket.consumeTokensAtomically([reservation], { + cooldownKeys: [], + deadlineAt: now + 5000, + signal: controller.signal, + }) + ) + if (interruption === 'deadline') vi.setSystemTime(now + 6000) + else controller.abort(new Error('cancelled')) + }) + if (interruption === 'deadline') { + expect(await pending?.result).toEqual({ allowed: false, retryAfterMs: 0 }) + } else { + await expect(pending?.result).rejects.toThrow('cancelled') + } + expect(await storedTokens()).toBe(10) + } + ) +}) From aa1ed38a03f8e8306958e0e3a9e4130c175be837 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 12 Oct 2026 14:53:51 -0700 Subject: [PATCH 4/5] fix(rate-limit): count only the lock wait in the atomic admission clock Sample clock_timestamp() before the first row can reach the lock and again after the last is held, so an uncontended admission decides at exactly the caller's clock and planning or execution time never shifts a refill. --- .../rate-limiter/storage/db-token-bucket.ts | 29 ++++++++++++------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts index d7adbd674e7..097931a60aa 100644 --- a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts +++ b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts @@ -42,11 +42,12 @@ function sqlArray(values: readonly (string | number)[], type: 'text' | 'float8') * Returns one row per reservation and per cooldown key, each carrying the decision, how many * buckets were spent, and the clock the decision used. * - * That clock is the caller's `sentAt` advanced by the time the statement spent acquiring its - * locks, in whole milliseconds: `clock` counts every `locked` row before reading - * `clock_timestamp()`, so the read happens only once the last lock is held. Measuring the wait - * on the database rather than reading its absolute time keeps every stored instant on the - * application clock that wrote it. + * That clock is the caller's `sentAt` advanced by the time spent acquiring the locks, in whole + * milliseconds. `lock_start` is joined into `locked`, so it is read before the first row can + * reach the lock; `clock` counts every `locked` row before reading `clock_timestamp()` again, so + * that read happens only once the last lock is held. An uncontended admission therefore decides + * at exactly `sentAt`, and measuring the wait rather than reading the database's absolute time + * keeps every stored instant on the application clock that wrote it. */ function buildAdmissionStatement( reservations: readonly TokenBucketReservation[], @@ -79,17 +80,23 @@ function buildAdmissionStatement( )} ) AS r(key, cost, max_tokens, refill_rate, refill_interval_ms) ), + lock_start AS ( + SELECT clock_timestamp() AS started_at + ), locked AS ( - SELECT key, tokens, last_refill_at, blocked_until - FROM ${rateLimitBucket} - WHERE key = ANY(${sqlArray(keys, 'text')}) - ORDER BY key - FOR UPDATE + SELECT bucket.key, bucket.tokens, bucket.last_refill_at, bucket.blocked_until + FROM ${rateLimitBucket} AS bucket + CROSS JOIN lock_start + WHERE bucket.key = ANY(${sqlArray(keys, 'text')}) + ORDER BY bucket.key + FOR UPDATE OF bucket ), clock AS ( SELECT ${sentAt}::float8 - + floor(extract(epoch FROM (held.locked_at - statement_timestamp())) * 1000) AS now_ms + + GREATEST(0, floor(extract(epoch FROM (held.locked_at - lock_start.started_at)) * 1000)) + AS now_ms FROM (SELECT clock_timestamp() AS locked_at FROM (SELECT count(*) FROM locked) every_lock) held + CROSS JOIN lock_start ), balance AS ( SELECT From ed2e401c1a2ab32f3fa51976b92b0e8f4f9dc271 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 12 Oct 2026 15:44:35 -0700 Subject: [PATCH 5/5] fix(connectors): keep counting every slice until a connector's first pass finishes A first pass's own unfinished slices wrote a count, so later slices of that pass saw it and stopped counting. Skip only once lastSyncAt shows a finished pass, which unfinished slices never advance. --- .../knowledge/connectors/sync-engine.test.ts | 21 ++++++++++++++----- .../lib/knowledge/connectors/sync-engine.ts | 20 +++++++++++------- 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index 29fe7aaf43b..c356acad943 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -1592,11 +1592,22 @@ describe('completeSuccessfulSync', () => { }) it.each([ - { recordedDocCount: 7, written: undefined }, - { recordedDocCount: null, written: 4 }, + { + scenario: 'after a finished pass', + lastSyncAt: new Date(), + lastSyncDocCount: 7, + written: undefined, + }, + { scenario: 'with no count yet', lastSyncAt: null, lastSyncDocCount: null, written: 4 }, + { + scenario: 'on a later slice of the first pass', + lastSyncAt: null, + lastSyncDocCount: 3, + written: 4, + }, ])( - 'leaves an unfinished pass on the last finished count unless none was recorded: $recordedDocCount', - async ({ recordedDocCount, written }) => { + 'counts an unfinished slice only until a pass has finished: $scenario', + async ({ lastSyncAt, lastSyncDocCount, written }) => { const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine') queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }]) queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1', lastSyncDocCount: 7 }]) @@ -1624,7 +1635,7 @@ describe('completeSuccessfulSync', () => { }, null, false, - recordedDocCount + { lastSyncAt, lastSyncDocCount } ) ).resolves.toBe(true) diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.ts b/apps/sim/lib/knowledge/connectors/sync-engine.ts index 34202895fe7..5102b99df22 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.ts @@ -400,8 +400,8 @@ export async function completeSuccessfulSync( directoryNotice: string | null = null, /** A pending ACL rewrite this run began but did not finish: the flag stays and the next run resumes it. */ accessRewriteUnfinished = false, - /** The count the connector held when this run took its lock. */ - recordedDocCount: number | null = null + /** The connector as this run locked it: whether a pass has finished, and the count it holds. */ + lockedConnector: { lastSyncAt: Date | null; lastSyncDocCount: number | null } | null = null ): Promise { const processingDispatchFailed = result.processingDispatch.failed > 0 const contentNotice = @@ -434,13 +434,17 @@ export async function completeSuccessfulSync( * adds and deletes would drift; when it cannot be read the previous count stands, which the * guard already treats as a floor. * - * A slice that leaves its content pass unfinished keeps the count the last finished pass - * recorded, the same pass `lastSyncAt` still names, so a large connector is not counted again - * on every slice of a pass whose count would only be replaced when it finishes. A connector - * with no count yet is still counted, so a first pass shows its progress. + * Once a pass has finished, a slice that leaves the next pass unfinished keeps the count that + * finished pass recorded, the same pass `lastSyncAt` still names, so a large connector is not + * counted again on every slice of a pass whose count is replaced when it finishes. Unfinished + * slices never advance `lastSyncAt`, so until a pass finishes every slice is counted and a + * first pass shows its progress. */ const keepRecordedCount = - contentPass !== undefined && !contentPass.complete && recordedDocCount !== null + contentPass !== undefined && + !contentPass.complete && + lockedConnector?.lastSyncAt != null && + lockedConnector.lastSyncDocCount !== null const actualDocCount = keepRecordedCount ? null : await countLiveConnectorDocuments(connectorId).catch((error: unknown) => { @@ -1429,7 +1433,7 @@ export async function executeSync( contentPass, directoryNotice, accessRewriteUnfinished, - connector.lastSyncDocCount + connector ) if (!completionLanded) {