From e87213e75692b4a203fd365dfe0b509122279194 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 11:25:42 -0700 Subject: [PATCH 1/5] feat(knowledge): fill the projection source and ACL in shards The fill's cost is graph insertion per row, which scales with cores while one run uses one. The runner takes an upper bound on its id range; a run may fill one shard of the chunk id space, sliced on the first hex digit so every shard is a contiguous range; the task admits one run per shard; the enqueue starts one chain per shard; and the warm after the fill runs once nothing is left unfilled anywhere. --- .../projection-source-acl-backfill.ts | 13 ++- .../projection-source-acl-backfill.test.ts | 85 +++++++++++++-- .../search/projection-source-acl-backfill.ts | 103 +++++++++++++++--- .../scripts/backfill-projection-source-acl.ts | 9 +- .../0021_embedding_search_connector.test.ts | 20 ++++ .../0021_embedding_search_connector.ts | 7 +- 6 files changed, 207 insertions(+), 30 deletions(-) diff --git a/apps/sim/background/projection-source-acl-backfill.ts b/apps/sim/background/projection-source-acl-backfill.ts index e6b75c0794a..133bb9f59e6 100644 --- a/apps/sim/background/projection-source-acl-backfill.ts +++ b/apps/sim/background/projection-source-acl-backfill.ts @@ -9,12 +9,19 @@ import { /** One run's share of the backfill, inside the worker's run ceiling with room to end its page. */ const RUN_BUDGET_MS = 60 * 60 * 1000 +/** + * Runs admitted at once: one per shard the id space may be sliced into. Shards fill disjoint + * ranges, so runs never fill the same page against each other; an unsharded chain still runs one + * at a time because each run triggers its continuation only as it ends. + */ +export const PROJECTION_SOURCE_ACL_BACKFILL_SHARDS = 4 + /** * 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, @@ -22,7 +29,7 @@ export const projectionSourceAclBackfillTask = task({ retry: { maxAttempts: 3 }, 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 }) 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..9c90e0a7c4e 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,13 +3,15 @@ */ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger } = vi.hoisted(() => ({ - mockBackfill: vi.fn(), - mockEnd: vi.fn(async () => undefined), - mockPostgres: vi.fn(), - mockPrewarm: vi.fn(async () => []), - mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })), -})) +const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger, mockUnsafe } = + vi.hoisted(() => ({ + mockBackfill: vi.fn(), + mockEnd: vi.fn(async () => undefined), + mockPostgres: vi.fn(), + mockPrewarm: vi.fn(async () => []), + mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })), + mockUnsafe: vi.fn(async () => [{ unfilled: false }]), + })) vi.mock('@sim/db', () => ({ resolveDbUrl: () => 'postgres://localhost:5432/sim' })) vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({ @@ -29,10 +31,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(() => { @@ -101,6 +104,56 @@ 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 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(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: 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', () => { @@ -118,7 +171,7 @@ 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'], }) expect(mockTasksTrigger).toHaveBeenCalledWith( 'projection-source-acl-backfill', @@ -127,4 +180,18 @@ describe('enqueueProjectionSourceAclBackfill', () => { ) expect(mockBackfill).not.toHaveBeenCalled() }) + + it('starts one run per shard, each on its own slice', async () => { + await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 }, 4)).resolves.toEqual({ + runIds: ['run-1', 'run-1', 'run-1', 'run-1'], + }) + expect(mockTasksTrigger.mock.calls.map(([, payload]) => payload)).toEqual( + [0, 1, 2, 3].map((index) => ({ pageSize: 25, shard: { index, count: 4 } })) + ) + }) + + it('refuses a shard count the id space cannot be sliced into before starting anything', async () => { + await expect(enqueueProjectionSourceAclBackfill({}, 3)).rejects.toThrow('must divide 16') + 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..a5ffe1b02fe 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,50 @@ 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 +} + 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 (!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 +79,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 +96,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,27 +113,61 @@ 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; put the ranking pages back before anyone + * searches. With shards, the one that finishes last does it: a shard that still finds unfilled + * rows anywhere leaves the warm to whoever fills them. Two shards ending in the same moment can + * both read none left and both warm, which repeats reads and nothing else. + */ + if (await projectionsFilled(sql)) { + await prewarmSearchProjection(sql, { budgetMs: PROJECTION_PREWARM_BUDGET_MS }) + } return null } finally { await sql.end() } } +/** Whether no projection still holds a row without its source and ACL; each read is one index probe. */ +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} WHERE acl IS NULL) AS unfilled` + ) + if (row?.unfilled) return false + } + return true +} + /** - * 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. + * 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. */ export async function enqueueProjectionSourceAclBackfill( - payload: ProjectionSourceAclBackfillPayload = {} -): Promise<{ runId: string }> { + payload: ProjectionSourceAclBackfillPayload = {}, + shards = 1 +): Promise<{ runIds: 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 } + const region = await resolveTriggerRegion() + if (shards !== 1) assertProjectionSourceAclShard({ index: 0, count: shards }) + const payloads: ProjectionSourceAclBackfillPayload[] = + shards === 1 + ? [payload] + : Array.from({ length: shards }, (_, index) => ({ + ...payload, + shard: { index, count: shards }, + })) + const runIds: string[] = [] + for (const shardPayload of payloads) { + const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, shardPayload, { + region, + }) + runIds.push(handle.id) + } + logger.info('Projection source and ACL backfill enqueued', { runIds, shards }) + return { runIds } } diff --git a/apps/sim/scripts/backfill-projection-source-acl.ts b/apps/sim/scripts/backfill-projection-source-acl.ts index be982412df6..01a66c316bd 100644 --- a/apps/sim/scripts/backfill-projection-source-acl.ts +++ b/apps/sim/scripts/backfill-projection-source-acl.ts @@ -22,10 +22,15 @@ import { const logger = createLogger('BackfillProjectionSourceAcl') -/** A script has no long-lived process to detach into, so without a worker it fills inline. */ +/** + * A script has no long-lived process to detach into, so without a worker it fills inline. With a + * worker, `--shards ` slices the id space so that many runs fill at once. + */ async function main(): Promise { if (isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) { - const handle = await enqueueProjectionSourceAclBackfill() + const flag = process.argv.indexOf('--shards') + const shards = flag === -1 ? 1 : Number(process.argv[flag + 1]) + const handle = await enqueueProjectionSourceAclBackfill({}, shards) logger.info('Backfill enqueued on the Trigger.dev worker', handle) return } 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..a7a048ced9c 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. */ @@ -135,6 +137,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 +152,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,7 +164,7 @@ 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 }) From 0de85331332ab13c00bd00e54dfe652830a5e3c0 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 11:42:31 -0700 Subject: [PATCH 2/5] feat(knowledge): make a sliced fill's start idempotent and leave analysis to whoever finishes A start's runs carry a per-shard idempotency key, so a retried start finds its runs rather than making more; a sliced start refuses a cursor, which belongs to one chain; the shard count is bounded by the runs the queue admits at once; a run reports its range done rather than the projection, and the run that finds nothing left anywhere analyzes and warms both projections once; the script reads its flag up front and refuses it where there is no worker to slice across. --- .../projection-source-acl-backfill.ts | 13 +++---- .../projection-source-acl-backfill.test.ts | 37 +++++++++++++++--- .../search/projection-source-acl-backfill.ts | 38 +++++++++++++++---- .../scripts/backfill-projection-source-acl.ts | 21 ++++++++-- .../0021_embedding_search_connector.ts | 11 ++++-- 5 files changed, 92 insertions(+), 28 deletions(-) diff --git a/apps/sim/background/projection-source-acl-backfill.ts b/apps/sim/background/projection-source-acl-backfill.ts index 133bb9f59e6..100eeef8777 100644 --- a/apps/sim/background/projection-source-acl-backfill.ts +++ b/apps/sim/background/projection-source-acl-backfill.ts @@ -1,6 +1,7 @@ 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, runProjectionSourceAclBackfill, @@ -9,13 +10,6 @@ import { /** One run's share of the backfill, inside the worker's run ceiling with room to end its page. */ const RUN_BUDGET_MS = 60 * 60 * 1000 -/** - * Runs admitted at once: one per shard the id space may be sliced into. Shards fill disjoint - * ranges, so runs never fill the same page against each other; an unsharded chain still runs one - * at a time because each run triggers its continuation only as it ends. - */ -export const PROJECTION_SOURCE_ACL_BACKFILL_SHARDS = 4 - /** * 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 @@ -27,6 +21,11 @@ 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: PROJECTION_SOURCE_ACL_BACKFILL_SHARDS, 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 9c90e0a7c4e..420040c02e1 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 @@ -63,8 +63,11 @@ 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({}) + 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( @@ -121,11 +124,12 @@ describe('runProjectionSourceAclBackfill', () => { expect(mockBackfill.mock.calls[1][2]).toMatchObject({ afterId: '4', beforeId: '8' }) }) - it('leaves the warm to whoever fills the rows another shard still holds', async () => { + 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) }) @@ -149,6 +153,7 @@ describe('projectionSourceAclShardRange', () => { 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) => { @@ -176,22 +181,42 @@ describe('enqueueProjectionSourceAclBackfill', () => { expect(mockTasksTrigger).toHaveBeenCalledWith( 'projection-source-acl-backfill', { pageSize: 25 }, - { region: 'us-east-1' } + { + region: 'us-east-1', + idempotencyKey: 'projection-source-acl-backfill:1:0', + idempotencyKeyTTL: '1h', + } ) expect(mockBackfill).not.toHaveBeenCalled() }) - it('starts one run per shard, each on its own slice', async () => { + it('starts one run per shard, each on its own slice with its own idempotency key', async () => { await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 }, 4)).resolves.toEqual({ runIds: ['run-1', 'run-1', 'run-1', 'run-1'], }) 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.idempotencyKey)).toEqual( + [0, 1, 2, 3].map((index) => `projection-source-acl-backfill:4:${index}`) + ) }) - it('refuses a shard count the id space cannot be sliced into before starting anything', async () => { - await expect(enqueueProjectionSourceAclBackfill({}, 3)).rejects.toThrow('must divide 16') + 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 a5ffe1b02fe..b7dd7142c10 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -36,6 +36,12 @@ export interface ProjectionSourceAclBackfillShard { 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 @@ -50,6 +56,11 @@ function assertProjectionSourceAclShard({ index, count }: ProjectionSourceAclBac 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}`) } @@ -117,12 +128,15 @@ export async function runProjectionSourceAclBackfill( elapsedMs: Date.now() - startedAt, }) /** - * The fill just streamed through both projections; put the ranking pages back before anyone - * searches. With shards, the one that finishes last does it: a shard that still finds unfilled - * rows anywhere leaves the warm to whoever fills them. Two shards ending in the same moment can - * both read none left and both warm, which repeats reads and nothing else. + * 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 @@ -142,18 +156,26 @@ async function projectionsFilled(sql: postgres.Sql): Promise { return true } +/** How long a start stays idempotent: long enough that a retried command finds its runs, not a second set. */ +const ENQUEUE_IDEMPOTENCY_TTL = '1h' + /** * 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. + * own slice at the same time. Safe to call again at any time: a run only fills rows still unset, + * and a start repeated within the hour finds the runs it already made rather than making more. 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 = {}, shards = 1 ): Promise<{ runIds: string[] }> { + if (shards !== 1) { + assertProjectionSourceAclShard({ index: 0, count: shards }) + if (payload.cursor) throw new Error('A sliced projection backfill cannot start from a cursor') + } const { tasks } = await import('@trigger.dev/sdk') const region = await resolveTriggerRegion() - if (shards !== 1) assertProjectionSourceAclShard({ index: 0, count: shards }) const payloads: ProjectionSourceAclBackfillPayload[] = shards === 1 ? [payload] @@ -162,9 +184,11 @@ export async function enqueueProjectionSourceAclBackfill( shard: { index, count: shards }, })) const runIds: string[] = [] - for (const shardPayload of payloads) { + for (const [index, shardPayload] of payloads.entries()) { const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, shardPayload, { region, + idempotencyKey: `${PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID}:${shards}:${index}`, + idempotencyKeyTTL: ENQUEUE_IDEMPOTENCY_TTL, }) runIds.push(handle.id) } diff --git a/apps/sim/scripts/backfill-projection-source-acl.ts b/apps/sim/scripts/backfill-projection-source-acl.ts index 01a66c316bd..0b61ea2792f 100644 --- a/apps/sim/scripts/backfill-projection-source-acl.ts +++ b/apps/sim/scripts/backfill-projection-source-acl.ts @@ -22,18 +22,31 @@ import { const logger = createLogger('BackfillProjectionSourceAcl') +/** `--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. With a - * worker, `--shards ` slices the id space so that many runs fill at once. + * 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 flag = process.argv.indexOf('--shards') - const shards = flag === -1 ? 1 : Number(process.argv[flag + 1]) const handle = await enqueueProjectionSourceAclBackfill({}, shards) logger.info('Backfill enqueued on the Trigger.dev worker', handle) 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.ts b/packages/db/script-migrations/0021_embedding_search_connector.ts index a7a048ced9c..3a05f4e76d3 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.ts @@ -114,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, @@ -170,8 +172,6 @@ export async function backfillProjectionSourceAcl( }) 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 @@ -192,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, From a53830cab3d2691b848faa8b53fba82fe899cbbc Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 11:55:27 -0700 Subject: [PATCH 3/5] feat(knowledge): find a fill chain in flight by its tag before starting another A start no longer relies on an idempotency key, which would have blocked a legitimate restart for its window and collided across options. Every run of a chain carries the chain's tag, continuations included, and a start lists in-flight runs by that tag first: a range whose chain is still running is left to it, whether the start is sliced or not, and a start after a chain ended starts anew. --- .../projection-source-acl-backfill.ts | 3 + .../projection-source-acl-backfill.test.ts | 67 ++++++++++++++----- .../search/projection-source-acl-backfill.ts | 49 +++++++++++--- .../scripts/backfill-projection-source-acl.ts | 4 +- 4 files changed, 92 insertions(+), 31 deletions(-) diff --git a/apps/sim/background/projection-source-acl-backfill.ts b/apps/sim/background/projection-source-acl-backfill.ts index 100eeef8777..71c80c86369 100644 --- a/apps/sim/background/projection-source-acl-backfill.ts +++ b/apps/sim/background/projection-source-acl-backfill.ts @@ -4,6 +4,7 @@ import { PROJECTION_SOURCE_ACL_BACKFILL_SHARDS, PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, type ProjectionSourceAclBackfillPayload, + projectionSourceAclChainTag, runProjectionSourceAclBackfill, } from '@/lib/knowledge/search/projection-source-acl-backfill' @@ -36,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 420040c02e1..590b7637998 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,15 +3,25 @@ */ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger, mockUnsafe } = - vi.hoisted(() => ({ - mockBackfill: vi.fn(), - mockEnd: vi.fn(async () => undefined), - mockPostgres: vi.fn(), - mockPrewarm: vi.fn(async () => []), - mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })), - mockUnsafe: vi.fn(async () => [{ unfilled: false }]), - })) +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 }> => (async function* () {})() + ), + mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })), + mockUnsafe: vi.fn(async () => [{ unfilled: false }]), +})) vi.mock('@sim/db', () => ({ resolveDbUrl: () => 'postgres://localhost:5432/sim' })) vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({ @@ -20,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) => { @@ -164,6 +177,8 @@ describe('projectionSourceAclShardRange', () => { 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', @@ -177,28 +192,44 @@ describe('enqueueProjectionSourceAclBackfill', () => { it('hands the backfill to the Trigger.dev worker when one is configured', async () => { await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 })).resolves.toEqual({ 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', - idempotencyKey: 'projection-source-acl-backfill:1:0', - idempotencyKeyTTL: '1h', - } + { region: 'us-east-1', tags: ['projection-source-acl-backfill:shard:0/1'] } ) expect(mockBackfill).not.toHaveBeenCalled() }) - it('starts one run per shard, each on its own slice with its own idempotency key', async () => { + 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' } + })() + ) + 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.idempotencyKey)).toEqual( - [0, 1, 2, 3].map((index) => `projection-source-acl-backfill:4:${index}`) + expect(mockTasksTrigger.mock.calls.map(([, , options]) => options.tags)).toEqual( + [0, 1, 2, 3].map((index) => [`projection-source-acl-backfill:shard:${index}/4`]) ) }) 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 b7dd7142c10..ba723d93832 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -156,25 +156,38 @@ async function projectionsFilled(sql: postgres.Sql): Promise { return true } -/** How long a start stays idempotent: long enough that a retried command finds its runs, not a second set. */ -const ENQUEUE_IDEMPOTENCY_TTL = '1h' +/** 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}` +} + +/** A run that has not ended: it, or the continuation it triggers, still owns its range. */ +const IN_FLIGHT_RUN_STATUSES = [ + 'PENDING_VERSION', + 'QUEUED', + 'DEQUEUED', + 'EXECUTING', + 'WAITING', + 'DELAYED', +] as const /** * 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 start repeated within the hour finds the runs it already made rather than making more. A - * cursor belongs to one chain, so a sliced start takes none: each slice begins at its own bound. + * 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 = {}, shards = 1 -): Promise<{ runIds: string[] }> { +): Promise<{ runIds: string[]; inFlight: string[] }> { if (shards !== 1) { assertProjectionSourceAclShard({ index: 0, count: shards }) if (payload.cursor) throw new Error('A sliced projection backfill cannot start from a cursor') } - const { tasks } = await import('@trigger.dev/sdk') + const { runs, tasks } = await import('@trigger.dev/sdk') const region = await resolveTriggerRegion() const payloads: ProjectionSourceAclBackfillPayload[] = shards === 1 @@ -184,14 +197,28 @@ export async function enqueueProjectionSourceAclBackfill( shard: { index, count: shards }, })) const runIds: string[] = [] - for (const [index, shardPayload] of payloads.entries()) { + const inFlight: string[] = [] + for (const shardPayload of payloads) { + const tag = projectionSourceAclChainTag(shardPayload.shard) + let running: string | undefined + for await (const run of runs.list({ + taskIdentifier: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, + tag, + status: [...IN_FLIGHT_RUN_STATUSES], + limit: 1, + })) { + running = run.id + } + if (running) { + inFlight.push(running) + continue + } const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, shardPayload, { region, - idempotencyKey: `${PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID}:${shards}:${index}`, - idempotencyKeyTTL: ENQUEUE_IDEMPOTENCY_TTL, + tags: [tag], }) runIds.push(handle.id) } - logger.info('Projection source and ACL backfill enqueued', { runIds, shards }) - return { runIds } + 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 0b61ea2792f..ebc4a1ec390 100644 --- a/apps/sim/scripts/backfill-projection-source-acl.ts +++ b/apps/sim/scripts/backfill-projection-source-acl.ts @@ -41,8 +41,8 @@ function shardsFlag(): number { async function main(): Promise { const shards = shardsFlag() if (isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) { - const handle = await enqueueProjectionSourceAclBackfill({}, shards) - 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) From 1cf2483e1e0b12ab01b1b2b314c82a9693979dfb Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 12:04:22 -0700 Subject: [PATCH 4/5] feat(knowledge): close the instant between a fill chain's lookup and its start The in-flight lookup and the trigger are two calls, so two starts in the same instant could both find no chain. Each trigger now carries the chain's tag as a short-lived idempotency key, long enough to cover that instant and short enough never to hold a later restart. --- .../search/projection-source-acl-backfill.test.ts | 7 ++++++- .../knowledge/search/projection-source-acl-backfill.ts | 9 +++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) 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 590b7637998..e400cd9cfe0 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 @@ -200,7 +200,12 @@ describe('enqueueProjectionSourceAclBackfill', () => { expect(mockTasksTrigger).toHaveBeenCalledWith( 'projection-source-acl-backfill', { pageSize: 25 }, - { region: 'us-east-1', tags: ['projection-source-acl-backfill:shard:0/1'] } + { + region: 'us-east-1', + tags: ['projection-source-acl-backfill:shard:0/1'], + idempotencyKey: 'projection-source-acl-backfill:shard:0/1', + idempotencyKeyTTL: '2m', + } ) expect(mockBackfill).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 ba723d93832..8ab790771a8 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -162,6 +162,13 @@ export function projectionSourceAclChainTag(shard?: ProjectionSourceAclBackfillS return `${PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID}:shard:${index}/${count}` } +/** + * How long a start's trigger stays idempotent. The in-flight lookup and the trigger are two + * calls, so two starts in the same instant could both find no chain; a key that lives just past + * that instant closes the gap without holding a later, legitimate restart. + */ +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 = [ 'PENDING_VERSION', @@ -216,6 +223,8 @@ export async function enqueueProjectionSourceAclBackfill( const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, shardPayload, { region, tags: [tag], + idempotencyKey: tag, + idempotencyKeyTTL: START_IDEMPOTENCY_TTL, }) runIds.push(handle.id) } From 41db9c386ebb3ed08c98fd9d2e99bd14461946df Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 12:26:54 -0700 Subject: [PATCH 5/5] feat(knowledge): key a fill start on the chain it follows, and count only rows the fill can finish A start's idempotency key is the chain's latest run, so two starts that saw the same state collapse into one while a start after another chain ended is its own; a start validates any shard it is handed; and the completion probe counts only unfilled rows whose document exists, since a row whose document is gone is not the fill's to finish. --- .../projection-source-acl-backfill.test.ts | 35 +++++++++++++++++-- .../search/projection-source-acl-backfill.ts | 34 +++++++++++------- 2 files changed, 53 insertions(+), 16 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 e400cd9cfe0..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 @@ -17,7 +17,7 @@ const { mockPostgres: vi.fn(), mockPrewarm: vi.fn(async () => []), mockRunsList: vi.fn( - (_query: unknown): AsyncIterable<{ id: string }> => (async function* () {})() + (_query: unknown): AsyncIterable<{ id: string; status: string }> => (async function* () {})() ), mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })), mockUnsafe: vi.fn(async () => [{ unfilled: false }]), @@ -78,6 +78,12 @@ 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') + ) + ).toBe(true) expect(mockUnsafe.mock.calls.map(([query]) => query)).toEqual( expect.arrayContaining(['ANALYZE embedding_search', 'ANALYZE embedding_keyword_tin']) ) @@ -203,17 +209,40 @@ describe('enqueueProjectionSourceAclBackfill', () => { { region: 'us-east-1', tags: ['projection-source-acl-backfill:shard:0/1'], - idempotencyKey: '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' } + if ((query as { tag: string }).tag.endsWith(':shard:1/4')) + yield { id: 'run-live', status: 'EXECUTING' } })() ) await expect(enqueueProjectionSourceAclBackfill({}, 4)).resolves.toEqual({ 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 8ab790771a8..4e9b0f0e813 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -145,11 +145,17 @@ export async function runProjectionSourceAclBackfill( } } -/** Whether no projection still holds a row without its source and ACL; each read is one index probe. */ +/** + * 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} WHERE acl IS NULL) AS unfilled` + `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 } @@ -163,21 +169,22 @@ export function projectionSourceAclChainTag(shard?: ProjectionSourceAclBackfillS } /** - * How long a start's trigger stays idempotent. The in-flight lookup and the trigger are two - * calls, so two starts in the same instant could both find no chain; a key that lives just past - * that instant closes the gap without holding a later, legitimate restart. + * 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 = [ +const IN_FLIGHT_RUN_STATUSES: ReadonlySet = new Set([ 'PENDING_VERSION', 'QUEUED', 'DEQUEUED', 'EXECUTING', 'WAITING', 'DELAYED', -] as const +]) /** * Starts the backfill on the deployment's Trigger.dev worker, where bounded runs chain until the @@ -190,6 +197,7 @@ export async function enqueueProjectionSourceAclBackfill( 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') @@ -207,23 +215,23 @@ export async function enqueueProjectionSourceAclBackfill( const inFlight: string[] = [] for (const shardPayload of payloads) { const tag = projectionSourceAclChainTag(shardPayload.shard) - let running: string | undefined + /** 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, - status: [...IN_FLIGHT_RUN_STATUSES], limit: 1, })) { - running = run.id + latest = run } - if (running) { - inFlight.push(running) + 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, + idempotencyKey: `${tag}:after:${latest?.id ?? 'none'}`, idempotencyKeyTTL: START_IDEMPOTENCY_TTL, }) runIds.push(handle.id)