fix(knowledge): retry a projection fill page after an index lock or statement timeout instead of ending the chain

Author: waleedlatif1

packages/db/script-migrations/0021_embedding_search_connector.test.ts

@@ -1,14 +1,114 @@
 /**
  * @vitest-environment node
  */
-import { backfillProjectionSourceAcl } from '@sim/db/script-migrations/0021_embedding_search_connector'
+import {
+  backfillProjectionSourceAcl,
+  PROJECTION_SOURCE_ACL_PAGE_RETRIES,
+} from '@sim/db/script-migrations/0021_embedding_search_connector'
 import type { Sql } from 'postgres'
-import { describe, expect, it, vi } from 'vitest'
+import { afterEach, describe, expect, it, vi } from 'vitest'
 
 /** A session that must never be reached: every case below is refused before the first page. */
 const untouched = { begin: vi.fn() } as unknown as Sql
 
+type PageRow = { scanned: number; filled: number; last_id: string | null }
+
+/** A driver error carrying a SQLSTATE, the shape `postgres` throws. */
+function postgresError(code: string): Error {
+  return Object.assign(new Error(`canceling statement (SQLSTATE ${code})`), { code })
+}
+
+/**
+ * A session whose page statement answers from `outcomes` in order — a row, or an error to throw —
+ * and records the cursor each page was bound to. `beforePage` runs with the page's index before
+ * it answers.
+ */
+function sessionOf(outcomes: Array<PageRow | Error>, beforePage?: (index: number) => void) {
+  const cursors: string[] = []
+  const tx = {
+    unsafe: vi.fn(async (query: string, params?: unknown[]) => {
+      if (!query.includes('WITH page')) return []
+      beforePage?.(cursors.length)
+      cursors.push(String(params?.[0]))
+      const outcome = outcomes.shift()
+      if (outcome === undefined) throw new Error('No outcome left for this page')
+      if (outcome instanceof Error) throw outcome
+      return [outcome]
+    }),
+  }
+  const session = {
+    begin: vi.fn(async (work: (tx: unknown) => Promise<unknown>) => work(tx)),
+    unsafe: vi.fn(async () => []),
+  } as unknown as Sql
+  return { session, cursors }
+}
+
+/** Runs a backfill under fake timers, so its retry pauses pass without waiting. */
+async function backfillNow(...args: Parameters<typeof backfillProjectionSourceAcl>) {
+  vi.useFakeTimers()
+  const result = backfillProjectionSourceAcl(...args)
+  /** A rejection must not surface as unhandled while the timers are still being drained. */
+  const settled = result.catch(() => undefined)
+  await vi.runAllTimersAsync()
+  await settled
+  return result
+}
+
 describe('backfillProjectionSourceAcl', () => {
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  it.each(['55P03', '57014'])(
+    'retries the page after a %s timeout and moves the cursor only once it commits',
+    async (code) => {
+      const { session, cursors } = sessionOf([
+        { scanned: 2, filled: 2, last_id: 'id-2' },
+        postgresError(code),
+        postgresError(code),
+        { scanned: 1, filled: 1, last_id: 'id-3' },
+        { scanned: 0, filled: 0, last_id: null },
+      ])
+      await expect(
+        backfillNow(session, 'embedding_keyword_tin', { pauseMs: 0 })
+      ).resolves.toMatchObject({ scanned: 3, written: 3, afterId: 'id-3', done: true })
+      expect(cursors).toEqual(['', 'id-2', 'id-2', 'id-2', 'id-3'])
+    }
+  )
+
+  it('gives up on a page that times out more than the retry limit in a row', async () => {
+    const { session, cursors } = sessionOf(
+      Array.from({ length: PROJECTION_SOURCE_ACL_PAGE_RETRIES + 1 }, () => postgresError('55P03'))
+    )
+    await expect(backfillNow(session, 'embedding_keyword_tin', { pauseMs: 0 })).rejects.toThrow(
+      'SQLSTATE 55P03'
+    )
+    expect(cursors).toHaveLength(PROJECTION_SOURCE_ACL_PAGE_RETRIES + 1)
+    expect(new Set(cursors)).toEqual(new Set(['']))
+  })
+
+  it('propagates an error that is not a timeout without retrying', async () => {
+    const { session, cursors } = sessionOf([postgresError('42P01')])
+    await expect(backfillNow(session, 'embedding_search', { pauseMs: 0 })).rejects.toThrow(
+      'SQLSTATE 42P01'
+    )
+    expect(cursors).toEqual([''])
+  })
+
+  it('leaves a page still failing at the budget to the continuation, from the last committed page', async () => {
+    const { session, cursors } = sessionOf(
+      [{ scanned: 1, filled: 1, last_id: 'id-1' }, postgresError('57014')],
+      /** The second page spends the budget before the database cancels it. */
+      (index) => {
+        if (index === 1) vi.advanceTimersByTime(1000)
+      }
+    )
+    await expect(
+      backfillNow(session, 'embedding_search', { pauseMs: 0, budgetMs: 1000 })
+    ).resolves.toMatchObject({ afterId: 'id-1', done: false })
+    expect(cursors).toEqual(['', 'id-1'])
+  })
+
   it.each([0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY])(
     'refuses a page size of %s instead of reporting the projection filled',
     async (pageSize) => {

packages/db/script-migrations/0021_embedding_search_connector.ts

@@ -1,5 +1,6 @@
 import { createLogger } from '@sim/logger'
 import { sleep } from '@sim/utils/helpers'
+import { backoffWithJitter } from '@sim/utils/retry'
 import postgres, { type Sql } from 'postgres'
 
 const logger = createLogger('ProjectionSourceAcl')
@@ -15,12 +16,37 @@ export const PROJECTION_SOURCE_ACL_PAGE_SIZE = 100
 export const PROJECTION_SOURCE_ACL_PAGE_PAUSE_MS = 250
 
 /**
- * Longest a page may run before the database cancels it; the run then fails and resumes. A caller
- * that bounds a run leaves at least this much headroom after its budget, since the budget is
- * checked between pages and the page in flight runs to this limit.
+ * Longest a page may run before the database cancels it; the page is then retried in place. A
+ * caller that bounds a run leaves at least this much headroom after its budget, plus the longest
+ * retry pause, since the budget is checked between pages and the page in flight runs to this limit.
  */
 export const PROJECTION_SOURCE_ACL_PAGE_TIMEOUT_MS = 60_000
 
+/**
+ * How many times in a row one page may time out before the run fails. With the pauses below a page
+ * waits out a few minutes of index maintenance before giving up.
+ */
+export const PROJECTION_SOURCE_ACL_PAGE_RETRIES = 6
+
+/** Pause before a page is retried: 10 s, doubling to 30 s, with jitter. */
+const PAGE_RETRY_PAUSE = { baseMs: 10_000, maxMs: 30_000 } as const
+
+/**
+ * The two ways the database cancels a page: `lock_timeout` (55P03) while the page's index write
+ * waits on a lock the index's background maintenance holds, and `statement_timeout` (57014) when
+ * the page itself runs past {@link PROJECTION_SOURCE_ACL_PAGE_TIMEOUT_MS}. Both pass once the
+ * maintenance moves on, so both are retried the same way.
+ */
+const PAGE_TIMEOUT_CODES: ReadonlySet<string> = new Set(['55P03', '57014'])
+
+/** The SQLSTATE on a driver error, or on the error it wraps. */
+function postgresErrorCode(error: unknown): string | undefined {
+  if (typeof error !== 'object' || error === null) return undefined
+  const code = (error as { code?: unknown }).code
+  if (typeof code === 'string') return code
+  return postgresErrorCode((error as { cause?: unknown }).cause)
+}
+
 /** Pages between progress log lines. */
 const PROGRESS_EVERY_PAGES = 100
 
@@ -117,6 +143,15 @@ export interface ProjectionSourceAclBackfillProgress {
  * predicate, the join per candidate every row paid before the columns existed. A run fills the
  * range it was given and reports that range done; whether the projection as a whole is done, and
  * the analysis the planner then needs, is the caller's, since several runs may share a projection.
+ *
+ * A page the database cancels — on a lock timeout, because the keyword index's background
+ * maintenance holds the index page the row's write needs, or on a statement timeout — is retried
+ * in place after a pause, up to {@link PROJECTION_SOURCE_ACL_PAGE_RETRIES} times in a row, and
+ * the cursor stays on the last page that committed. A retry from the run's original cursor would
+ * instead walk every row the run had filled, past the index entries those writes left behind,
+ * into a statement timeout of its own; and the maintenance that cancelled the page outlasts the
+ * few attempts a run gets, so the fill would end where it stalled. A page that is still failing
+ * when the budget runs out is left to the continuation rather than retried past it.
  */
 export async function backfillProjectionSourceAcl(
   sql: Sql,
@@ -144,32 +179,54 @@ export async function backfillProjectionSourceAcl(
   let written = 0
   let pages = 0
   let done = false
+  /** Timeouts in a row on the page after `afterId`; reset once it commits. */
+  let timeouts = 0
   for (;;) {
-    const page = await sql.begin(async (tx) => {
-      await tx.unsafe("SET LOCAL lock_timeout = '5s'")
-      await tx.unsafe(`SET LOCAL statement_timeout = ${PROJECTION_SOURCE_ACL_PAGE_TIMEOUT_MS}`)
-      const [row] = await tx.unsafe<
-        Array<{ scanned: number; filled: number; last_id: string | null }>
-      >(
-        `WITH page AS (
-          SELECT s.id, s.document_id, d.connector_id, d.acl
-          FROM ${projection} s JOIN document d ON d.id = s.document_id
-          WHERE s.id > $1 AND ($2::text IS NULL OR s.id < $2) AND s.acl IS NULL
-          ORDER BY s.id LIMIT ${pageSize}
-          FOR SHARE OF d
-        ), updated AS (
-          UPDATE ${projection} s SET connector_id = page.connector_id, acl = page.acl
-          FROM page
-          WHERE s.id = page.id AND s.document_id = page.document_id AND s.acl IS NULL
-          RETURNING s.id
+    let page: { scanned: number; filled: number; last_id: string | null }
+    try {
+      page = await sql.begin(async (tx) => {
+        await tx.unsafe("SET LOCAL lock_timeout = '5s'")
+        await tx.unsafe(`SET LOCAL statement_timeout = ${PROJECTION_SOURCE_ACL_PAGE_TIMEOUT_MS}`)
+        const [row] = await tx.unsafe<
+          Array<{ scanned: number; filled: number; last_id: string | null }>
+        >(
+          `WITH page AS (
+            SELECT s.id, s.document_id, d.connector_id, d.acl
+            FROM ${projection} s JOIN document d ON d.id = s.document_id
+            WHERE s.id > $1 AND ($2::text IS NULL OR s.id < $2) AND s.acl IS NULL
+            ORDER BY s.id LIMIT ${pageSize}
+            FOR SHARE OF d
+          ), updated AS (
+            UPDATE ${projection} s SET connector_id = page.connector_id, acl = page.acl
+            FROM page
+            WHERE s.id = page.id AND s.document_id = page.document_id AND s.acl IS NULL
+            RETURNING s.id
+          )
+          SELECT (SELECT count(*)::int FROM page) AS scanned,
+            (SELECT count(*)::int FROM updated) AS filled,
+            (SELECT max(id) FROM page) AS last_id`,
+          [afterId, beforeId]
         )
-        SELECT (SELECT count(*)::int FROM page) AS scanned,
-          (SELECT count(*)::int FROM updated) AS filled,
-          (SELECT max(id) FROM page) AS last_id`,
-        [afterId, beforeId]
-      )
-      return row
-    })
+        return row
+      })
+    } catch (error) {
+      const code = postgresErrorCode(error)
+      if (code === undefined || !PAGE_TIMEOUT_CODES.has(code)) throw error
+      timeouts += 1
+      if (timeouts > PROJECTION_SOURCE_ACL_PAGE_RETRIES) throw error
+      if (Date.now() >= deadline) break
+      const pauseMs = backoffWithJitter(timeouts, null, PAGE_RETRY_PAUSE)
+      logger.warn('Projection source and ACL backfill page timed out; retrying', {
+        projection,
+        afterId,
+        code,
+        attempt: timeouts,
+        retryInMs: Math.round(pauseMs),
+      })
+      await sleep(pauseMs)
+      continue
+    }
+    timeouts = 0
     if (page.last_id === null) {
       done = true
       break