fix(knowledge): stop the fill page retry at the run budget and only on statement timeouts

Author: waleedlatif1

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

@@ -13,9 +13,15 @@ const untouched = { begin: vi.fn() } as unknown as Sql
 
 type PageRow = { scanned: number; filled: number; last_id: string | null }
 
+/** The message the database pairs with each cancellation SQLSTATE. */
+const CANCELLATION_MESSAGES: Record<string, string> = {
+  '55P03': 'canceling statement due to lock timeout',
+  '57014': 'canceling statement due to statement timeout',
+}
+
 /** 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 })
+function postgresError(code: string, message = CANCELLATION_MESSAGES[code] ?? 'failed'): Error {
+  return Object.assign(new Error(`${message} (SQLSTATE ${code})`), { code })
 }
 
 /**
@@ -95,6 +101,27 @@ describe('backfillProjectionSourceAcl', () => {
     expect(cursors).toEqual([''])
   })
 
+  it('propagates an explicit cancellation, which shares the statement timeout SQLSTATE', async () => {
+    const { session, cursors } = sessionOf([
+      postgresError('57014', 'canceling statement due to user request'),
+    ])
+    await expect(backfillNow(session, 'embedding_search', { pauseMs: 0 })).rejects.toThrow(
+      'user request'
+    )
+    expect(cursors).toEqual([''])
+  })
+
+  it('does not start another page when the budget ran out during the retry pause', async () => {
+    const { session, cursors } = sessionOf([
+      { scanned: 1, filled: 1, last_id: 'id-1' },
+      postgresError('57014'),
+    ])
+    await expect(
+      backfillNow(session, 'embedding_search', { pauseMs: 0, budgetMs: 1000 })
+    ).resolves.toMatchObject({ afterId: 'id-1', done: false })
+    expect(cursors).toEqual(['', 'id-1'])
+  })
+
   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')],

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

@@ -25,21 +25,13 @@ 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. Index maintenance was
  * observed holding the page for several minutes; with the pauses below a page waits roughly
- * twelve minutes, so the budget covers one such pass.
+ * twelve minutes, about fourteen at the jitter's worst, so the budget covers one such pass.
  */
 export const PROJECTION_SOURCE_ACL_PAGE_RETRIES = 12
 
-/** Pause before a page is retried: 10 s, doubling to 60 s, with jitter. */
+/** Pause before a page is retried: 10 s, doubling to a 60 s base with up to 20% jitter (about 72 s). */
 const PAGE_RETRY_PAUSE = { baseMs: 10_000, maxMs: 60_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
@@ -48,6 +40,28 @@ function postgresErrorCode(error: unknown): string | undefined {
   return postgresErrorCode((error as { cause?: unknown }).cause)
 }
 
+/** The message on a driver error, or on the error it wraps. */
+function postgresErrorMessage(error: unknown): string | undefined {
+  if (typeof error !== 'object' || error === null) return undefined
+  const message = (error as { message?: unknown }).message
+  if (typeof message === 'string') return message
+  return postgresErrorMessage((error as { cause?: unknown }).cause)
+}
+
+/**
+ * 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. 57014 is also what an explicit
+ * cancellation raises, and that is not retried: only the message tells the two apart.
+ */
+function isPageTimeout(error: unknown): boolean {
+  const code = postgresErrorCode(error)
+  if (code === '55P03') return true
+  if (code !== '57014') return false
+  return postgresErrorMessage(error)?.includes('statement timeout') ?? false
+}
+
 /** Pages between progress log lines. */
 const PROGRESS_EVERY_PAGES = 100
 
@@ -211,8 +225,8 @@ export async function backfillProjectionSourceAcl(
         return row
       })
     } catch (error) {
+      if (!isPageTimeout(error)) throw 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
@@ -225,6 +239,8 @@ export async function backfillProjectionSourceAcl(
         retryInMs: Math.round(pauseMs),
       })
       await sleep(pauseMs)
+      /** Checked again after the pause, so a timeout at the budget cannot start another page. */
+      if (Date.now() >= deadline) break
       continue
     }
     timeouts = 0