diff --git a/apps/sim/background/projection-source-acl-backfill.ts b/apps/sim/background/projection-source-acl-backfill.ts index e6b75c0794a..71c80c86369 100644 --- a/apps/sim/background/projection-source-acl-backfill.ts +++ b/apps/sim/background/projection-source-acl-backfill.ts @@ -1,8 +1,10 @@ import { task, tasks } from '@trigger.dev/sdk' import { resolveTriggerRegion } from '@/lib/core/async-jobs/region' import { + PROJECTION_SOURCE_ACL_BACKFILL_SHARDS, PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, type ProjectionSourceAclBackfillPayload, + projectionSourceAclChainTag, runProjectionSourceAclBackfill, } from '@/lib/knowledge/search/projection-source-acl-backfill' @@ -13,16 +15,21 @@ const RUN_BUDGET_MS = 60 * 60 * 1000 * Trigger.dev wrapper around `runProjectionSourceAclBackfill`. A run fills unset rows for up to * {@link RUN_BUDGET_MS}, then triggers its continuation from the cursor it reached, so the whole * projection is filled across as many bounded runs as it takes. Retry-safe: every run writes only - * rows still unset, so a retried or restarted run repeats no write. The queue admits one run at a - * time, so two starts never fill the same pages against each other. + * rows still unset, so a retried or restarted run repeats no write. A shard's continuation keeps + * its shard, so a sliced fill stays sliced until every slice is done. */ export const projectionSourceAclBackfillTask = task({ id: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, machine: 'small-1x', retry: { maxAttempts: 3 }, + /** + * One run per shard the id space may be sliced into. Shards fill disjoint ranges, so runs never + * fill the same page against each other; an unsliced chain still runs one at a time because each + * run triggers its continuation only as it ends. + */ queue: { name: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, - concurrencyLimit: 1, + concurrencyLimit: PROJECTION_SOURCE_ACL_BACKFILL_SHARDS, }, run: async (payload: ProjectionSourceAclBackfillPayload) => { const cursor = await runProjectionSourceAclBackfill(payload, { budgetMs: RUN_BUDGET_MS }) @@ -30,6 +37,8 @@ export const projectionSourceAclBackfillTask = task({ const continuation: ProjectionSourceAclBackfillPayload = { ...payload, cursor } await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, continuation, { region: await resolveTriggerRegion(), + /** The chain's tag rides on every continuation, so a start finds the chain wherever it is. */ + tags: [projectionSourceAclChainTag(payload.shard)], }) }, }) 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 4a357e3e9a3..f5d9434e2a1 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 @@ -3,12 +3,24 @@ */ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger } = vi.hoisted(() => ({ +const { + mockBackfill, + mockEnd, + mockPostgres, + mockPrewarm, + mockRunsList, + mockTasksTrigger, + mockUnsafe, +} = vi.hoisted(() => ({ mockBackfill: vi.fn(), mockEnd: vi.fn(async () => undefined), mockPostgres: vi.fn(), mockPrewarm: vi.fn(async () => []), + mockRunsList: vi.fn( + (_query: unknown): AsyncIterable<{ id: string; status: string }> => (async function* () {})() + ), mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })), + mockUnsafe: vi.fn(async () => [{ unfilled: false }]), })) vi.mock('@sim/db', () => ({ resolveDbUrl: () => 'postgres://localhost:5432/sim' })) @@ -18,7 +30,10 @@ vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({ })) vi.mock('postgres', () => ({ default: mockPostgres })) vi.mock('@/lib/knowledge/search/prewarm', () => ({ prewarmSearchProjection: mockPrewarm })) -vi.mock('@trigger.dev/sdk', () => ({ tasks: { trigger: mockTasksTrigger } })) +vi.mock('@trigger.dev/sdk', () => ({ + runs: { list: mockRunsList }, + tasks: { trigger: mockTasksTrigger }, +})) vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' })) vi.mock('@/lib/core/utils/background', () => ({ runDetached: (_label: string, work: () => Promise) => { @@ -29,10 +44,11 @@ vi.mock('@/lib/core/utils/background', () => ({ import { enqueueProjectionSourceAclBackfill, PROJECTION_PREWARM_BUDGET_MS, + projectionSourceAclShardRange, runProjectionSourceAclBackfill, } from '@/lib/knowledge/search/projection-source-acl-backfill' -const connection = { end: mockEnd } +const connection = { end: mockEnd, unsafe: mockUnsafe } describe('runProjectionSourceAclBackfill', () => { beforeEach(() => { @@ -60,8 +76,17 @@ describe('runProjectionSourceAclBackfill', () => { expect(mockEnd).toHaveBeenCalledTimes(1) }) - it('warms the projections on the same connection once both are filled, before closing it', async () => { + 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') + ) + ).toBe(true) + expect(mockUnsafe.mock.calls.map(([query]) => query)).toEqual( + expect.arrayContaining(['ANALYZE embedding_search', 'ANALYZE embedding_keyword_tin']) + ) expect(mockPrewarm).toHaveBeenCalledTimes(1) expect(mockPrewarm).toHaveBeenCalledWith(connection, { budgetMs: PROJECTION_PREWARM_BUDGET_MS }) expect(mockPrewarm.mock.invocationCallOrder[0]).toBeLessThan( @@ -101,11 +126,65 @@ describe('runProjectionSourceAclBackfill', () => { await expect(runProjectionSourceAclBackfill({})).rejects.toThrow('statement timeout') expect(mockEnd).toHaveBeenCalledTimes(1) }) + + it('fills only its shard of the id space in both projections', async () => { + await runProjectionSourceAclBackfill({ shard: { index: 1, count: 4 } }) + for (const [, , options] of mockBackfill.mock.calls) { + expect(options).toMatchObject({ afterId: '4', beforeId: '8' }) + } + }) + + it('resumes a shard after its cursor and keeps its upper bound', async () => { + await runProjectionSourceAclBackfill({ + shard: { index: 1, count: 4 }, + cursor: { projection: 'embedding_search', afterId: '5a' }, + }) + expect(mockBackfill.mock.calls[0][2]).toMatchObject({ afterId: '5a', beforeId: '8' }) + expect(mockBackfill.mock.calls[1][2]).toMatchObject({ afterId: '4', beforeId: '8' }) + }) + + it('leaves the analysis and the warm to whoever fills the rows another shard still holds', async () => { + mockUnsafe.mockResolvedValueOnce([{ unfilled: true }]) + await expect( + runProjectionSourceAclBackfill({ shard: { index: 0, count: 4 } }) + ).resolves.toBeNull() + expect(mockUnsafe.mock.calls.map(([query]) => query)).not.toContain('ANALYZE embedding_search') + expect(mockPrewarm).not.toHaveBeenCalled() + expect(mockEnd).toHaveBeenCalledTimes(1) + }) +}) + +describe('projectionSourceAclShardRange', () => { + it('slices the hex id space into contiguous ranges', () => { + expect(projectionSourceAclShardRange({ index: 0, count: 4 })).toEqual({ + afterId: '', + beforeId: '4', + }) + expect(projectionSourceAclShardRange({ index: 3, count: 4 })).toEqual({ + afterId: 'c', + beforeId: undefined, + }) + expect(projectionSourceAclShardRange({ index: 0, count: 1 })).toEqual({ + afterId: '', + beforeId: undefined, + }) + }) + + it.each([ + [{ index: 0, count: 3 }, 'shard count must divide 16'], + [{ index: 0, count: 8 }, 'shard count must be at most 4'], + [{ index: 4, count: 4 }, 'shard index must be within 0..3'], + [{ index: 0.5, count: 2 }, 'shard index must be within 0..1'], + ])('refuses %j', (shard, message) => { + expect(() => projectionSourceAclShardRange(shard)).toThrow(message) + }) }) describe('enqueueProjectionSourceAclBackfill', () => { beforeEach(() => { vi.clearAllMocks() + /** No chain in flight unless a case says so. */ + mockRunsList.mockImplementation(() => (async function* () {})()) mockPostgres.mockReturnValue(connection) mockBackfill.mockResolvedValue({ projection: 'embedding_search', @@ -118,13 +197,91 @@ describe('enqueueProjectionSourceAclBackfill', () => { it('hands the backfill to the Trigger.dev worker when one is configured', async () => { await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 })).resolves.toEqual({ - runId: 'run-1', + runIds: ['run-1'], + inFlight: [], }) + expect(mockRunsList).toHaveBeenCalledWith( + expect.objectContaining({ tag: 'projection-source-acl-backfill:shard:0/1' }) + ) expect(mockTasksTrigger).toHaveBeenCalledWith( 'projection-source-acl-backfill', { pageSize: 25 }, - { region: 'us-east-1' } + { + region: 'us-east-1', + tags: ['projection-source-acl-backfill:shard:0/1'], + idempotencyKey: 'projection-source-acl-backfill:shard:0/1:after:none', + idempotencyKeyTTL: '2m', + } ) expect(mockBackfill).not.toHaveBeenCalled() }) + + it('keys a start after a chain that ended on that chain, so a restart is its own start', async () => { + mockRunsList.mockImplementation(() => + (async function* () { + yield { id: 'run-done', status: 'COMPLETED' } + })() + ) + await expect(enqueueProjectionSourceAclBackfill({})).resolves.toEqual({ + runIds: ['run-1'], + inFlight: [], + }) + expect(mockTasksTrigger.mock.calls[0][2].idempotencyKey).toBe( + 'projection-source-acl-backfill:shard:0/1:after:run-done' + ) + }) + + it('refuses a shard the id space cannot be sliced into before starting anything', async () => { + await expect( + enqueueProjectionSourceAclBackfill({ shard: { index: 5, count: 4 } }) + ).rejects.toThrow('shard index must be within 0..3') + expect(mockTasksTrigger).not.toHaveBeenCalled() + }) + + it('leaves a range whose chain is still in flight to that chain', async () => { + mockRunsList.mockImplementation((query: unknown) => + (async function* () { + if ((query as { tag: string }).tag.endsWith(':shard:1/4')) + yield { id: 'run-live', status: 'EXECUTING' } + })() + ) + await expect(enqueueProjectionSourceAclBackfill({}, 4)).resolves.toEqual({ + runIds: ['run-1', 'run-1', 'run-1'], + inFlight: ['run-live'], + }) + expect(mockTasksTrigger.mock.calls.map(([, payload]) => payload.shard?.index)).toEqual([ + 0, 2, 3, + ]) + }) + + it('starts one run per shard, each on its own slice under its own chain tag', async () => { + await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 }, 4)).resolves.toEqual({ + runIds: ['run-1', 'run-1', 'run-1', 'run-1'], + inFlight: [], + }) + expect(mockTasksTrigger.mock.calls.map(([, payload]) => payload)).toEqual( + [0, 1, 2, 3].map((index) => ({ pageSize: 25, shard: { index, count: 4 } })) + ) + expect(mockTasksTrigger.mock.calls.map(([, , options]) => options.tags)).toEqual( + [0, 1, 2, 3].map((index) => [`projection-source-acl-backfill:shard:${index}/4`]) + ) + }) + + it.each([ + [3, 'must divide 16'], + [8, 'must be at most 4'], + ])('refuses %s shards before starting anything', async (shards, message) => { + await expect(enqueueProjectionSourceAclBackfill({}, shards)).rejects.toThrow(message) + expect(mockTasksTrigger).not.toHaveBeenCalled() + }) + + it('refuses to slice a start that carries a cursor, which belongs to one chain', async () => { + await expect( + enqueueProjectionSourceAclBackfill( + { cursor: { projection: 'embedding_search', afterId: '5a' } }, + 4 + ) + ).rejects.toThrow('cannot start from a cursor') + expect(mockTasksTrigger).not.toHaveBeenCalled() + }) }) 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 6ac3947d5d3..4e9b0f0e813 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -26,13 +26,61 @@ export interface ProjectionSourceAclBackfillCursor { afterId: string } +/** + * One of `count` equal slices of the chunk id space. Chunk ids are lowercase hex UUIDs, so the + * space is sliced on the first hex digit and `count` must divide sixteen; every slice is a + * contiguous id range, so workers on different slices never fill the same page. + */ +export interface ProjectionSourceAclBackfillShard { + index: number + count: number +} + +/** + * Shards the id space may be sliced into at most: the runs the task's queue admits at once, so a + * sliced fill runs its slices together rather than in waves. + */ +export const PROJECTION_SOURCE_ACL_BACKFILL_SHARDS = 4 + export interface ProjectionSourceAclBackfillPayload { /** The first projection's first page when absent. */ cursor?: ProjectionSourceAclBackfillCursor + /** The whole id space when absent. */ + shard?: ProjectionSourceAclBackfillShard pageSize?: number pauseMs?: number } +/** Refuses a shard the id space cannot be sliced into. */ +function assertProjectionSourceAclShard({ index, count }: ProjectionSourceAclBackfillShard): void { + if (!Number.isInteger(count) || count < 1 || 16 % count !== 0) { + throw new Error(`Projection backfill shard count must divide 16, got ${count}`) + } + if (count > PROJECTION_SOURCE_ACL_BACKFILL_SHARDS) { + throw new Error( + `Projection backfill shard count must be at most ${PROJECTION_SOURCE_ACL_BACKFILL_SHARDS}, got ${count}` + ) + } + if (!Number.isInteger(index) || index < 0 || index >= count) { + throw new Error(`Projection backfill shard index must be within 0..${count - 1}, got ${index}`) + } +} + +/** The id range a shard covers: the id its first page follows, and the first id past it. */ +export function projectionSourceAclShardRange(shard: ProjectionSourceAclBackfillShard): { + afterId: string + beforeId: string | undefined +} { + assertProjectionSourceAclShard(shard) + const { index, count } = shard + const width = 16 / count + const digit = (value: number) => value.toString(16) + return { + afterId: index === 0 ? '' : digit(index * width), + beforeId: index + 1 < count ? digit((index + 1) * width) : undefined, + } +} + export interface ProjectionSourceAclBackfillRunOptions { /** Stop once this much time has passed and return where to resume; unbounded otherwise. */ budgetMs?: number @@ -42,8 +90,9 @@ export interface ProjectionSourceAclBackfillRunOptions { * Fills the ranking projections' source and ACL columns from the cursor onwards, one projection * after the other, on a connection of its own: the page statement binds the keyset cursor as a * scalar and needs no array parameter, so the pool's options would serve, but a run this long - * should not hold one of the worker's pooled connections. Returns the cursor to continue from when - * the budget ran out, `null` once both projections are filled. + * should not hold one of the worker's pooled connections. A shard fills its own slice of the id + * space in every projection. Returns the cursor to continue from when the budget ran out, `null` + * once the run's range is filled in both projections. */ export async function runProjectionSourceAclBackfill( payload: ProjectionSourceAclBackfillPayload, @@ -58,13 +107,16 @@ export async function runProjectionSourceAclBackfill( ? PROJECTION_SOURCE_ACL_TABLES.indexOf(payload.cursor.projection) : 0 if (start < 0) throw new Error(`Unknown projection ${payload.cursor?.projection}`) + const range = payload.shard ? projectionSourceAclShardRange(payload.shard) : undefined for (const projection of PROJECTION_SOURCE_ACL_TABLES.slice(start)) { const budgetMs = options.budgetMs === undefined ? undefined : Math.max(0, options.budgetMs - (Date.now() - startedAt)) const progress = await backfillProjectionSourceAcl(sql, projection, { - afterId: payload.cursor?.projection === projection ? payload.cursor.afterId : undefined, + afterId: + payload.cursor?.projection === projection ? payload.cursor.afterId : range?.afterId, + beforeId: range?.beforeId, pageSize: payload.pageSize, pauseMs: payload.pauseMs, budgetMs, @@ -72,10 +124,21 @@ export async function runProjectionSourceAclBackfill( if (!progress.done) return { projection, afterId: progress.afterId } } logger.info('Projection source and ACL backfill complete', { + shard: payload.shard, elapsedMs: Date.now() - startedAt, }) - /** The fill just streamed through both projections; put the ranking pages back before anyone searches. */ - await prewarmSearchProjection(sql, { budgetMs: PROJECTION_PREWARM_BUDGET_MS }) + /** + * The fill just streamed through both projections: the planner last saw every row unfilled and + * should see the finished projections, and the ranking pages should be back before anyone + * searches. With shards, the one that finishes last does both: a shard that still finds + * unfilled rows anywhere leaves them to whoever fills those. Two shards ending in the same + * moment can both read none left and both do this, which repeats reads and nothing else. + */ + if (await projectionsFilled(sql)) { + for (const projection of PROJECTION_SOURCE_ACL_TABLES) + await sql.unsafe(`ANALYZE ${projection}`) + await prewarmSearchProjection(sql, { budgetMs: PROJECTION_PREWARM_BUDGET_MS }) + } return null } finally { await sql.end() @@ -83,16 +146,96 @@ export async function runProjectionSourceAclBackfill( } /** - * Starts the backfill on the deployment's Trigger.dev worker, where bounded runs chain until both - * projections are filled. Safe to call again at any time: a run only fills rows still unset. + * 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. + */ +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` + ) + if (row?.unfilled) return false + } + return true +} + +/** The tag every run of one chain carries, so a chain in flight is found before another is started. */ +export function projectionSourceAclChainTag(shard?: ProjectionSourceAclBackfillShard): string { + const { index, count } = shard ?? { index: 0, count: 1 } + return `${PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID}:shard:${index}/${count}` +} + +/** + * How long a start's trigger stays idempotent. The lookup and the trigger are two calls, so two + * starts in the same instant could both find no chain in flight; the key is what they both saw, + * the chain's latest run, so they collapse into one start, while a start after another chain has + * ended sees a different latest run and is a new key. + */ +const START_IDEMPOTENCY_TTL = '2m' + +/** A run that has not ended: it, or the continuation it triggers, still owns its range. */ +const IN_FLIGHT_RUN_STATUSES: ReadonlySet = new Set([ + 'PENDING_VERSION', + 'QUEUED', + 'DEQUEUED', + 'EXECUTING', + 'WAITING', + 'DELAYED', +]) + +/** + * Starts the backfill on the deployment's Trigger.dev worker, where bounded runs chain until the + * projections are filled: one chain over the whole id space, or one per shard, each filling its + * own slice at the same time. Safe to call again at any time: a run only fills rows still unset, + * and a range whose chain is still in flight is left to that chain rather than given a second. + * A cursor belongs to one chain, so a sliced start takes none: each slice begins at its own bound. */ export async function enqueueProjectionSourceAclBackfill( - payload: ProjectionSourceAclBackfillPayload = {} -): Promise<{ runId: string }> { - const { tasks } = await import('@trigger.dev/sdk') - const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, payload, { - region: await resolveTriggerRegion(), - }) - logger.info('Projection source and ACL backfill enqueued', { runId: handle.id }) - return { runId: handle.id } + payload: ProjectionSourceAclBackfillPayload = {}, + shards = 1 +): Promise<{ runIds: string[]; inFlight: string[] }> { + if (payload.shard) assertProjectionSourceAclShard(payload.shard) + if (shards !== 1) { + assertProjectionSourceAclShard({ index: 0, count: shards }) + if (payload.cursor) throw new Error('A sliced projection backfill cannot start from a cursor') + } + const { runs, tasks } = await import('@trigger.dev/sdk') + const region = await resolveTriggerRegion() + const payloads: ProjectionSourceAclBackfillPayload[] = + shards === 1 + ? [payload] + : Array.from({ length: shards }, (_, index) => ({ + ...payload, + shard: { index, count: shards }, + })) + const runIds: string[] = [] + const inFlight: string[] = [] + for (const shardPayload of payloads) { + const tag = projectionSourceAclChainTag(shardPayload.shard) + /** The chain's latest run, newest first, whatever its state. */ + let latest: { id: string; status: string } | undefined + for await (const run of runs.list({ + taskIdentifier: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, + tag, + limit: 1, + })) { + latest = run + } + if (latest && IN_FLIGHT_RUN_STATUSES.has(latest.status)) { + inFlight.push(latest.id) + continue + } + const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, shardPayload, { + region, + tags: [tag], + idempotencyKey: `${tag}:after:${latest?.id ?? 'none'}`, + idempotencyKeyTTL: START_IDEMPOTENCY_TTL, + }) + runIds.push(handle.id) + } + logger.info('Projection source and ACL backfill enqueued', { runIds, inFlight, shards }) + return { runIds, inFlight } } diff --git a/apps/sim/scripts/backfill-projection-source-acl.ts b/apps/sim/scripts/backfill-projection-source-acl.ts index be982412df6..ebc4a1ec390 100644 --- a/apps/sim/scripts/backfill-projection-source-acl.ts +++ b/apps/sim/scripts/backfill-projection-source-acl.ts @@ -22,13 +22,31 @@ import { const logger = createLogger('BackfillProjectionSourceAcl') -/** A script has no long-lived process to detach into, so without a worker it fills inline. */ +/** `--shards ` as given, or one; a value that is not a whole number is refused here. */ +function shardsFlag(): number { + const flag = process.argv.indexOf('--shards') + if (flag === -1) return 1 + const shards = Number(process.argv[flag + 1]) + if (!Number.isInteger(shards) || shards < 1) { + throw new Error(`--shards must be a positive whole number, got ${process.argv[flag + 1]}`) + } + return shards +} + +/** + * A script has no long-lived process to detach into, so without a worker it fills inline, one + * pass over the whole id space. With a worker, `--shards ` slices the id space so that many + * runs fill at once; inline there is nothing to slice across, so the flag is refused there. + */ async function main(): Promise { + const shards = shardsFlag() if (isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) { - const handle = await enqueueProjectionSourceAclBackfill() - logger.info('Backfill enqueued on the Trigger.dev worker', handle) + const started = await enqueueProjectionSourceAclBackfill({}, shards) + logger.info('Backfill enqueued on the Trigger.dev worker', started) return } + if (shards !== 1) + throw new Error('--shards needs the Trigger.dev worker; the inline fill is one pass') await runProjectionSourceAclBackfill({}) logger.info('Backfill complete') } 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 d4e030b449a..94c4246100e 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.test.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.test.ts @@ -25,4 +25,24 @@ describe('backfillProjectionSourceAcl', () => { ).rejects.toThrow('pause must be a non-negative number') expect(untouched.begin).not.toHaveBeenCalled() }) + + it('binds the range it was given to every page, so shards never meet', async () => { + const statements: Array<{ query: string; params?: unknown[] }> = [] + const tx = { + unsafe: vi.fn(async (query: string, params?: unknown[]) => { + statements.push({ query, params }) + return query.includes('WITH page') ? [{ scanned: 0, filled: 0, last_id: null }] : [] + }), + } + const session = { + begin: vi.fn(async (work: (tx: unknown) => Promise) => work(tx)), + unsafe: vi.fn(async () => []), + } as unknown as Sql + await expect( + backfillProjectionSourceAcl(session, 'embedding_search', { afterId: '4', beforeId: '8' }) + ).resolves.toMatchObject({ done: true }) + const page = statements.find((statement) => statement.query.includes('WITH page'))! + expect(page.query).toContain('($2::text IS NULL OR s.id < $2)') + expect(page.params).toEqual(['4', '8']) + }) }) diff --git a/packages/db/script-migrations/0021_embedding_search_connector.ts b/packages/db/script-migrations/0021_embedding_search_connector.ts index 4fc87e10fe3..3a05f4e76d3 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.ts @@ -83,6 +83,8 @@ export async function installProjectionSourceAcl(sql: Sql): Promise { export interface ProjectionSourceAclBackfillOptions { /** Resume after this chunk id; the projection's first page otherwise. */ afterId?: string + /** Stop before this chunk id; the projection's end otherwise. Lets workers fill disjoint ranges. */ + beforeId?: string pageSize?: number pauseMs?: number /** Stop once this much time has passed and report where to resume; unbounded otherwise. */ @@ -112,7 +114,9 @@ export interface ProjectionSourceAclBackfillProgress { * cost of a page and far more than a deploy can wait for; the run paces itself with a pause between * pages and stops at its budget so a background task can chain runs until the projection is filled. * Search does not wait: an unfilled row is decided on its document by the on-row candidate - * predicate, the join per candidate every row paid before the columns existed. + * 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. */ export async function backfillProjectionSourceAcl( sql: Sql, @@ -135,6 +139,7 @@ export async function backfillProjectionSourceAcl( const deadline = options.budgetMs === undefined ? Number.POSITIVE_INFINITY : startedAt + options.budgetMs let afterId = options.afterId ?? '' + const beforeId = options.beforeId ?? null let scanned = 0 let written = 0 let pages = 0 @@ -149,7 +154,7 @@ export async function backfillProjectionSourceAcl( `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 s.acl IS NULL + 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 ( @@ -161,14 +166,12 @@ export async function backfillProjectionSourceAcl( SELECT (SELECT count(*)::int FROM page) AS scanned, (SELECT count(*)::int FROM updated) AS filled, (SELECT max(id) FROM page) AS last_id`, - [afterId] + [afterId, beforeId] ) return row }) if (page.last_id === null) { done = true - /** The planner last saw every row unfilled; it should see the finished projection. */ - await sql.unsafe(`ANALYZE ${projection}`) break } afterId = page.last_id @@ -189,9 +192,12 @@ export async function backfillProjectionSourceAcl( if (Date.now() >= deadline) break } logger.info( - done ? 'Projection source and ACL backfilled' : 'Projection source and ACL backfill paused', + done + ? 'Projection source and ACL range backfilled' + : 'Projection source and ACL backfill paused', { projection, + beforeId, scanned, written, afterId,