improvement(db): cut hot-path connector count and rate-limit lock costs (#8944)

* 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.

* 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.

* 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.

* 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.

* 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.

Author: waleedlatif1

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 })

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([])
 

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,191 @@ 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<string, unknown> {
+  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
+  now_ms: 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 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 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[],
+  cooldownKeys: readonly string[],
+  keys: readonly string[],
+  sentAt: 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)
+    ),
+    lock_start AS (
+      SELECT clock_timestamp() AS started_at
+    ),
+    locked AS (
+      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
+        + 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
+        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 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((clock.now_ms - 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, 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' + 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, 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, 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,
+      o.now_ms
+    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 +263,31 @@ 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. 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[],
     options: AtomicAdmissionOptions
   ): Promise<AtomicAdmissionResult> {
     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 +302,20 @@ 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 sentAt = Date.now()
+        const rows = await tx.execute<AdmissionRow>(
+          buildAdmissionStatement(reservations, options.cooldownKeys, keys, sentAt)
         )
-        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 }
+        const now = rows[0]?.now_ms ?? sentAt
+        options.signal?.throwIfAborted()
+        if (Math.max(now, 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<Date | null> {