From 01398e075bbdf2b82005c63bc954af000400a9a2 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 06:34:53 -0700 Subject: [PATCH 1/3] fix(knowledge): retry a projection fill page after an index lock or statement timeout instead of ending the chain --- .../0021_embedding_search_connector.test.ts | 104 +++++++++++++++- .../0021_embedding_search_connector.ts | 111 +++++++++++++----- 2 files changed, 186 insertions(+), 29 deletions(-) diff --git a/packages/db/script-migrations/0021_embedding_search_connector.test.ts b/packages/db/script-migrations/0021_embedding_search_connector.test.ts index 94c4246100e..f15029a5ca7 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.test.ts +++ b/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, 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) => 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) { + 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) => { diff --git a/packages/db/script-migrations/0021_embedding_search_connector.ts b/packages/db/script-migrations/0021_embedding_search_connector.ts index 3a05f4e76d3..339248786bd 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.ts +++ b/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 = 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 From 7ea6231c76a5a04c9bfea0e6d220995c101e2c96 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 06:36:47 -0700 Subject: [PATCH 2/3] fix(knowledge): size the fill page retry budget to one index maintenance pass --- .../0021_embedding_search_connector.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/packages/db/script-migrations/0021_embedding_search_connector.ts b/packages/db/script-migrations/0021_embedding_search_connector.ts index 339248786bd..a2a46731d60 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.ts @@ -23,13 +23,14 @@ export const PROJECTION_SOURCE_ACL_PAGE_PAUSE_MS = 250 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. + * 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. */ -export const PROJECTION_SOURCE_ACL_PAGE_RETRIES = 6 +export const PROJECTION_SOURCE_ACL_PAGE_RETRIES = 12 -/** 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 +/** Pause before a page is retried: 10 s, doubling to 60 s, with jitter. */ +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 From 4c96292446b589978ffa23e28590b2b2109c2aa9 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 06:51:52 -0700 Subject: [PATCH 3/3] fix(knowledge): stop the fill page retry at the run budget and only on statement timeouts --- .../0021_embedding_search_connector.test.ts | 31 ++++++++++++++- .../0021_embedding_search_connector.ts | 38 +++++++++++++------ 2 files changed, 56 insertions(+), 13 deletions(-) diff --git a/packages/db/script-migrations/0021_embedding_search_connector.test.ts b/packages/db/script-migrations/0021_embedding_search_connector.test.ts index f15029a5ca7..d2218247e7a 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.test.ts +++ b/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 = { + '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')], diff --git a/packages/db/script-migrations/0021_embedding_search_connector.ts b/packages/db/script-migrations/0021_embedding_search_connector.ts index a2a46731d60..d615b3235da 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.ts +++ b/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 = 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