From 06d0f7dfd0ff85cdc21d2055ff81e96fc292f36c Mon Sep 17 00:00:00 2001 From: Waleed Date: Tue, 22 Sep 2026 06:57:13 -0700 Subject: [PATCH 1/2] fix(knowledge): retry a projection fill page after an index lock or statement timeout instead of ending the chain (#8137) * fix(knowledge): retry a projection fill page after an index lock or statement timeout instead of ending the chain * fix(knowledge): size the fill page retry budget to one index maintenance pass * fix(knowledge): stop the fill page retry at the run budget and only on statement timeouts --- .../0021_embedding_search_connector.test.ts | 131 +++++++++++++++++- .../0021_embedding_search_connector.ts | 128 +++++++++++++---- 2 files changed, 230 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..d2218247e7a 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,141 @@ /** * @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 } + +/** 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, message = CANCELLATION_MESSAGES[code] ?? 'failed'): Error { + return Object.assign(new Error(`${message} (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('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')], + /** 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..d615b3235da 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,52 @@ 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. Index maintenance was + * observed holding the page for several minutes; with the pauses below a page waits roughly + * 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 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 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) +} + +/** 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 @@ -117,6 +158,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 +194,56 @@ 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) { + if (!isPageTimeout(error)) throw error + const code = postgresErrorCode(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) + /** Checked again after the pause, so a timeout at the budget cannot start another page. */ + if (Date.now() >= deadline) break + continue + } + timeouts = 0 if (page.last_id === null) { done = true break From 17e415d1c210df47ebc6284c4abade80feebf17c Mon Sep 17 00:00:00 2001 From: Waleed Date: Tue, 22 Sep 2026 09:39:42 -0700 Subject: [PATCH 2/2] fix(knowledge): read the projection-filled probe off the partial index instead of a sequential scan (#8139) * fix(knowledge): read the projection-filled probe off the partial index instead of a sequential scan * test(knowledge): pin the backfill probe's ordered, capped shape --- .../projection-source-acl-backfill.test.ts | 19 ++++++++++++++----- .../search/projection-source-acl-backfill.ts | 11 +++++++---- apps/sim/lib/knowledge/search/queries.test.ts | 18 ++++++++++++++++++ apps/sim/lib/knowledge/search/queries.ts | 13 +++++++++++-- 4 files changed, 50 insertions(+), 11 deletions(-) diff --git a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts index f5d9434e2a1..4e0bb9a7d46 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts @@ -78,12 +78,21 @@ describe('runProjectionSourceAclBackfill', () => { it('analyzes and warms the projections on the same connection once both are filled, before closing it', async () => { await runProjectionSourceAclBackfill({}) - /** A row whose document is gone is not the fill's to finish; the probe joins the document. */ - expect( - mockUnsafe.mock.calls.some(([query]) => - String(query).includes('JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL') + /** + * A row whose document is gone is not the fill's to finish; the probe joins the document. It is + * ordered and capped so only the unfilled-rows index can serve it: an `EXISTS` drops both and + * leaves the planner a sequential scan of the projection. + */ + const probes = mockUnsafe.mock.calls + .map(([query]) => String(query).replace(/\s+/g, ' ')) + .filter((query) => query.includes('AS unfilled')) + expect(probes).toHaveLength(2) + for (const probe of probes) { + expect(probe).not.toContain('EXISTS') + expect(probe).toContain( + 'JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL ORDER BY s.id DESC LIMIT 1 ) IS NOT NULL AS unfilled' ) - ).toBe(true) + } expect(mockUnsafe.mock.calls.map(([query]) => query)).toEqual( expect.arrayContaining(['ANALYZE embedding_search', 'ANALYZE embedding_keyword_tin']) ) diff --git a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts index 4e9b0f0e813..0e6162d246b 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -148,14 +148,17 @@ export async function runProjectionSourceAclBackfill( /** * Whether no projection still holds a row the fill could give its source and ACL: a row without * them whose document exists. A row whose document is gone is not the fill's to finish and never - * counts as left. Each read is one index probe while any such row remains. + * counts as left. Each read is one probe of the unfilled-rows index while any such row remains: + * ordered by id and capped at one row so the planner cannot take a sequential scan, which an + * `EXISTS` would leave open by dropping the order and the limit. */ async function projectionsFilled(sql: postgres.Sql): Promise { for (const projection of PROJECTION_SOURCE_ACL_TABLES) { const [row] = await sql.unsafe>( - `SELECT EXISTS ( - SELECT 1 FROM ${projection} s JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL - ) AS unfilled` + `SELECT ( + SELECT s.id FROM ${projection} s JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL + ORDER BY s.id DESC LIMIT 1 + ) IS NOT NULL AS unfilled` ) if (row?.unfilled) return false } diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 219f323d72e..75fb7644604 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -2822,6 +2822,24 @@ describe('filters on a resolved scope', () => { expect(caps.at(-1)).toBe('20000') }) + it('looks for an unfilled row through the ordered, capped read the partial index serves', async () => { + traversedRows = [{ id: 'a' }] + rerankRows = [hit('a', 'src-a')] + queueTableRows(schemaMock.embedding, rerankRows) + await handleVectorOnlySearch({ + ...params, + permitted: { kind: 'unbounded', broad: true }, + accessPlan: plan(), + }) + const probes = statements().filter((query) => query.sql.includes('AS unfilled')) + expect(probes).toHaveLength(1) + /** An `EXISTS` drops its order and limit, and the planner then takes a sequential scan. */ + expect(probes[0].sql).not.toContain('EXISTS') + expect(probes[0].sql.replace(/\s+/g, ' ')).toContain( + 'SELECT ( SELECT ? FROM ? WHERE ? IS NULL ORDER BY ? DESC LIMIT 1 ) IS NOT NULL AS unfilled' + ) + }) + it('tests the date through the document inside an on-row walk when the filtered set is unbounded', async () => { traversedRows = [{ id: 'a' }] rerankRows = [hit('a', 'src-a')] diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index f2ff5bf066c..129f0082ac2 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -115,7 +115,13 @@ const PROJECTION_FILLED_TTL_MS = 60_000 /** * Whether the ranking projection still holds rows the backfill has not filled. Read off the - * unfilled-rows index in microseconds and remembered briefly: the answer only ever changes once. + * unfilled-rows index in milliseconds and remembered briefly: the answer only ever changes once. + * + * The read asks for the last unfilled row by id, not whether one exists: an `EXISTS` drops its + * order and limit, and while most rows are unfilled the planner expects a sequential scan to + * meet one at once, then walks the whole projection when the unfilled rows sit past the filled + * ones. Ordered by id and capped at one row, the read can only be the partial index, whose + * last entry is the row the fill reaches last. */ const projectionFilled = new LRUCache< ProjectionSourceAclTable, @@ -134,7 +140,10 @@ const projectionFilled = new LRUCache< try { const [row] = await runSearchQuery(context.budget, context.stage, (executor) => executor.execute<{ unfilled: boolean }>(sql` - SELECT EXISTS (SELECT 1 FROM ${table} WHERE ${table.acl} IS NULL) AS unfilled`) + SELECT ( + SELECT ${table.id} FROM ${table} WHERE ${table.acl} IS NULL + ORDER BY ${table.id} DESC LIMIT 1 + ) IS NOT NULL AS unfilled`) ) return !row?.unfilled } catch {