From 89b2bcd9113b3eef8d8a57e39dbf3b01de7a36c5 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Thu, 17 Sep 2026 13:08:11 -0700 Subject: [PATCH] fix(file-search): bound dispatcher claims and skip locked rows --- .../search/dispatcher.integration.ts | 70 ++++++++++++++++++ .../lib/workspace-files/search/dispatcher.ts | 72 ++++++++----------- 2 files changed, 100 insertions(+), 42 deletions(-) diff --git a/apps/sim/lib/workspace-files/search/dispatcher.integration.ts b/apps/sim/lib/workspace-files/search/dispatcher.integration.ts index 2a9eadf2599..afc6eee134c 100644 --- a/apps/sim/lib/workspace-files/search/dispatcher.integration.ts +++ b/apps/sim/lib/workspace-files/search/dispatcher.integration.ts @@ -61,6 +61,11 @@ describe('workspace file search dispatch PostgreSQL deadlines', () => { workspace_id text PRIMARY KEY, enqueued_at timestamp NOT NULL, updated_at timestamp NOT NULL, last_dispatched_at timestamp )` + await connection`CREATE INDEX ON workspace_file_search_index + (workspace_id, updated_at, file_id, source_content_updated_at) + WHERE status = 'pending' AND dispatched_at IS NULL` + await connection`CREATE INDEX ON workspace_file_search_index (workspace_id, dispatched_at) + WHERE status = 'pending' AND dispatched_at IS NOT NULL` await connection`INSERT INTO workspace_file_search_backfill (id, updated_at) VALUES ('workspace-file-search-v1', '2026-09-16 00:00:00')` database.current = drizzle(connection) @@ -92,6 +97,71 @@ describe('workspace file search dispatch PostgreSQL deadlines', () => { }) } + async function seedQueue(workspaceId: string, queued: number, active = 0) { + await connection`UPDATE workspace_file_search_backfill SET completed_at = now()` + await connection`INSERT INTO workspace_files (id, workspace_id, context, content_updated_at) + SELECT ${workspaceId} || '-' || lpad(n::text, 6, '0'), ${workspaceId}, 'workspace', '2026-09-16' + FROM generate_series(1, ${queued + active}) n` + await connection`INSERT INTO workspace_file_search_index + (file_id, workspace_id, source_content_updated_at, status, updated_at, dispatched_at) + SELECT id, workspace_id, content_updated_at, 'pending', '2026-09-16', + CASE WHEN row_number() OVER (ORDER BY id DESC) <= ${active} THEN now() ELSE NULL END + FROM workspace_files WHERE workspace_id = ${workspaceId}` + await connection`INSERT INTO workspace_file_search_dispatch_queue + (workspace_id, enqueued_at, updated_at) VALUES (${workspaceId}, now(), now())` + } + + it('skips a locked candidate without losing it or exceeding workspace capacity', async () => { + await seedQueue('workspace-1', 3, 1) + await connection.begin(async (tx) => { + await tx`SELECT file_id FROM workspace_file_search_index + WHERE file_id = 'workspace-1-000001' FOR UPDATE` + const results = await Promise.all([ + prepareWorkspaceFileSearchDispatch(), + prepareWorkspaceFileSearchDispatch(), + ]) + expect(results.flatMap((result) => result.payloads).map((payload) => payload.fileId)).toEqual( + ['workspace-1-000002'] + ) + const [locked] = await tx`SELECT dispatched_at FROM workspace_file_search_index + WHERE file_id = 'workspace-1-000001'` + expect(locked.dispatched_at).toBeNull() + }) + expect((await prepareWorkspaceFileSearchDispatch()).payloads).toEqual([]) + await connection`UPDATE workspace_file_search_index SET status = 'ready' + WHERE file_id = 'workspace-1-000002'` + const retry = await prepareWorkspaceFileSearchDispatch() + expect(retry.payloads.map((payload) => payload.fileId)).toEqual(['workspace-1-000001']) + }) + + it('claims only available slots from a large backlog and preserves current-file eligibility', async () => { + await seedQueue('workspace-1', 10_000, 1) + await seedQueue('workspace-2', 3) + await connection`UPDATE workspace_files SET deleted_at = now() WHERE id = 'workspace-1-000001'` + await connection`UPDATE workspace_files SET content_updated_at = '2026-09-17' + WHERE id = 'workspace-1-000002'` + await connection`UPDATE workspace_files SET context = 'execution' WHERE id = 'workspace-1-000003'` + const result = await prepareWorkspaceFileSearchDispatch() + expect(result.payloads.map((payload) => payload.fileId).sort()).toEqual([ + 'workspace-1-000004', + 'workspace-2-000001', + 'workspace-2-000002', + ]) + expect((await prepareWorkspaceFileSearchDispatch()).payloads).toEqual([]) + }) + + it('honors the remaining global capacity across workspace probes', async () => { + await seedQueue('workspace-1', 3) + await seedQueue('workspace-2', 3) + await seedQueue('workspace-active', 0, 99) + const result = await prepareWorkspaceFileSearchDispatch() + expect(result.payloads).toHaveLength(1) + expect((await prepareWorkspaceFileSearchDispatch()).payloads).toEqual([]) + const [row] = await connection`SELECT count(*)::int AS active FROM workspace_file_search_index + WHERE status = 'pending' AND dispatched_at IS NOT NULL` + expect(row.active).toBe(100) + }) + it('fails on a locked backfill row and releases the dispatcher lock', async () => { let release = () => {} let locked = () => {} diff --git a/apps/sim/lib/workspace-files/search/dispatcher.ts b/apps/sim/lib/workspace-files/search/dispatcher.ts index 91f4b6d1899..a9346754a62 100644 --- a/apps/sim/lib/workspace-files/search/dispatcher.ts +++ b/apps/sim/lib/workspace-files/search/dispatcher.ts @@ -301,6 +301,7 @@ async function reapStaleClaims(tx: DbTransaction, now: Date): Promise { return rows.length } +/** Probe each workspace's available slots and lock candidates before the update, skipping busy rows. */ async function claimQueuedWorkspaceJobs( tx: DbTransaction, workspaceIds: readonly string[], @@ -320,48 +321,35 @@ async function claimQueuedWorkspaceJobs( WITH selected_workspace(workspace_id) AS ( VALUES ${workspaceValues} ), - workspace_active AS ( - SELECT search_index.workspace_id, count(*)::int AS active_count - FROM workspace_file_search_index AS search_index - INNER JOIN selected_workspace AS selected - ON selected.workspace_id = search_index.workspace_id - WHERE search_index.status = 'pending' - AND search_index.dispatched_at IS NOT NULL - GROUP BY search_index.workspace_id - ), - ranked AS ( - SELECT - search_index.workspace_id, - search_index.file_id, - search_index.source_content_updated_at, - search_index.updated_at, - coalesce(workspace_active.active_count, 0) AS active_count, - row_number() OVER ( - PARTITION BY search_index.workspace_id - ORDER BY - search_index.updated_at, - search_index.file_id, - search_index.source_content_updated_at - ) AS workspace_rank - FROM workspace_file_search_index AS search_index - INNER JOIN selected_workspace AS selected - ON selected.workspace_id = search_index.workspace_id - INNER JOIN workspace_files AS file - ON file.id = search_index.file_id - AND file.workspace_id = search_index.workspace_id - AND file.context = 'workspace' - AND file.deleted_at IS NULL - AND file.content_updated_at = search_index.source_content_updated_at - LEFT JOIN workspace_active - ON workspace_active.workspace_id = search_index.workspace_id - WHERE search_index.status = 'pending' - AND search_index.dispatched_at IS NULL - ), - candidates AS ( - SELECT workspace_id, file_id, source_content_updated_at - FROM ranked - WHERE workspace_rank <= ${FILE_SEARCH_INDEX_WORKSPACE_OUTSTANDING} - active_count - ORDER BY updated_at, workspace_id, file_id, source_content_updated_at + candidates AS MATERIALIZED ( + SELECT queued.* + FROM selected_workspace AS selected + CROSS JOIN LATERAL ( + SELECT count(*)::int AS active_count + FROM ( + SELECT 1 FROM workspace_file_search_index AS active + WHERE active.workspace_id = selected.workspace_id + AND active.status = 'pending' AND active.dispatched_at IS NOT NULL + LIMIT ${FILE_SEARCH_INDEX_WORKSPACE_OUTSTANDING} + ) AS active_claims + ) AS workspace_active + CROSS JOIN LATERAL ( + SELECT search_index.workspace_id, search_index.file_id, + search_index.source_content_updated_at, search_index.updated_at + FROM workspace_file_search_index AS search_index + INNER JOIN workspace_files AS file + ON file.id = search_index.file_id + AND file.workspace_id = search_index.workspace_id + AND file.context = 'workspace' + AND file.deleted_at IS NULL + AND file.content_updated_at = search_index.source_content_updated_at + WHERE search_index.workspace_id = selected.workspace_id + AND search_index.status = 'pending' AND search_index.dispatched_at IS NULL + ORDER BY search_index.updated_at, search_index.file_id, search_index.source_content_updated_at + LIMIT greatest(0, ${FILE_SEARCH_INDEX_WORKSPACE_OUTSTANDING} - workspace_active.active_count) + FOR UPDATE OF search_index SKIP LOCKED + ) AS queued + ORDER BY queued.updated_at, queued.workspace_id, queued.file_id, queued.source_content_updated_at LIMIT ${remainingGlobalCapacity} ) UPDATE workspace_file_search_index AS search_index