From 69b10e3f1f01d5e50a3232b104e691136242bb84 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 17 Sep 2026 23:54:53 -0700 Subject: [PATCH 1/2] fix(file-search): stop cleanup starting a batch it cannot fund The batch loop admitted another batch whenever any budget remained, then passed that remainder through as the batch's statement timeout, clamped to 1ms. A batch admitted with a sliver left either runs past the budget it was given or aborts on its own statement timeout, which the caller reports as a cleanup failure rather than as work still to do. Stop once less than one batch's nominal share of the budget remains. The floor is derived from the budget and the batch cap rather than fixed, so retuning either cannot leave it admitting batches at more than their share again. --- .../search/chunks.integration.ts | 23 +++++++++++++++++++ .../lib/workspace-files/search/constants.ts | 10 ++++++++ .../lib/workspace-files/search/index-state.ts | 9 ++++---- 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/apps/sim/lib/workspace-files/search/chunks.integration.ts b/apps/sim/lib/workspace-files/search/chunks.integration.ts index ee6825ff29c..c957b18e8a2 100644 --- a/apps/sim/lib/workspace-files/search/chunks.integration.ts +++ b/apps/sim/lib/workspace-files/search/chunks.integration.ts @@ -25,6 +25,7 @@ vi.mock('@/lib/file-parsers', () => ({ parseBuffer: vi.fn(), isSupportedFileType import { FILE_SEARCH_CLEANUP_BATCH_ROWS, + FILE_SEARCH_CLEANUP_BUDGET_MS, FILE_SEARCH_CLEANUP_MAX_BATCHES, } from '@/lib/workspace-files/search/constants' import { prepareWorkspaceFileSearchDispatch } from '@/lib/workspace-files/search/dispatcher' @@ -373,6 +374,28 @@ describe('chunked workspace file search on PostgreSQL', () => { (await connection`SELECT count(*)::int AS count FROM workspace_file_search_chunk`)[0].count ).toBe(0) }) + it('ends a cleanup run cleanly when its time budget runs out', async () => { + const build = (await beginFileSearchBuild(revision))! + await connection`INSERT INTO workspace_file_search_chunk (build_id, workspace_id, ordinal, line_start, fragment, content) + SELECT ${build.id}, 'workspace-1', n, n + 1, false, 'x' FROM generate_series(0, 999) n` + await connection`UPDATE workspace_file_search_build SET expires_at = now() WHERE id = ${build.id}` + + /** The deadline is read first; every read after it reports a budget all but consumed. */ + const startedAt = Date.now() + const clock = vi + .spyOn(Date, 'now') + .mockReturnValueOnce(startedAt) + .mockReturnValue(startedAt + FILE_SEARCH_CLEANUP_BUDGET_MS - 1) + try { + await expect(cleanupFileSearchBuilds()).resolves.toBe(0) + } finally { + clock.mockRestore() + } + + expect( + (await connection`SELECT count(*)::int AS count FROM workspace_file_search_chunk`)[0].count + ).toBe(1000) + }) it('retires many small builds within one cleanup run', async () => { await connection`INSERT INTO workspace_file_search_build (id, file_id, workspace_id, source_content_updated_at, expires_at) SELECT 'retired-' || n, 'file-1', 'workspace-1', now(), now() FROM generate_series(1, 100) n` diff --git a/apps/sim/lib/workspace-files/search/constants.ts b/apps/sim/lib/workspace-files/search/constants.ts index 8cff0a3fb7b..bf4214fd636 100644 --- a/apps/sim/lib/workspace-files/search/constants.ts +++ b/apps/sim/lib/workspace-files/search/constants.ts @@ -32,6 +32,16 @@ export const FILE_SEARCH_CLEANUP_BATCH_BUILDS = 100 export const FILE_SEARCH_CLEANUP_BACKLOG_ROWS = 10000 export const FILE_SEARCH_CLEANUP_MAX_BATCHES = 10 export const FILE_SEARCH_CLEANUP_BUDGET_MS = 5000 +/** + * One batch's nominal share of the run budget, and so the smallest slice worth starting another + * with. + * + * Running out of budget is how cleanup normally ends. A batch admitted with less than its share + * either runs past the budget it was given or aborts on its own statement timeout, which the caller + * reports as a cleanup failure rather than as work still to do. + */ +export const FILE_SEARCH_CLEANUP_MIN_BATCH_MS = + FILE_SEARCH_CLEANUP_BUDGET_MS / FILE_SEARCH_CLEANUP_MAX_BATCHES export const FILE_SEARCH_RECONCILE_INTERVAL_MS = 60 * 60 * 1000 export const FILE_SEARCH_INSERT_BATCH_ROWS = 250 export const FILE_SEARCH_INSERT_BATCH_BYTES = 1024 * 1024 diff --git a/apps/sim/lib/workspace-files/search/index-state.ts b/apps/sim/lib/workspace-files/search/index-state.ts index a7d14520406..98a748efa72 100644 --- a/apps/sim/lib/workspace-files/search/index-state.ts +++ b/apps/sim/lib/workspace-files/search/index-state.ts @@ -14,6 +14,7 @@ import { FILE_SEARCH_CLEANUP_BATCH_ROWS, FILE_SEARCH_CLEANUP_BUDGET_MS, FILE_SEARCH_CLEANUP_MAX_BATCHES, + FILE_SEARCH_CLEANUP_MIN_BATCH_MS, FILE_SEARCH_INSERT_BATCH_BYTES, FILE_SEARCH_INSERT_BATCH_ROWS, } from '@/lib/workspace-files/search/constants' @@ -257,11 +258,11 @@ export async function failFileSearchRevision( export async function cleanupFileSearchBuilds(): Promise { const deadline = Date.now() + FILE_SEARCH_CLEANUP_BUDGET_MS let deleted = 0 - for (let batch = 0; batch < FILE_SEARCH_CLEANUP_MAX_BATCHES && Date.now() < deadline; batch++) { + for (let batch = 0; batch < FILE_SEARCH_CLEANUP_MAX_BATCHES; batch++) { + const remainingBudget = deadline - Date.now() + if (remainingBudget < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) break const result = await db.transaction(async (tx) => { - await configureFileSearchTransaction(tx, { - statementTimeout: Math.max(1, deadline - Date.now()), - }) + await configureFileSearchTransaction(tx, { statementTimeout: remainingBudget }) const builds = await tx.execute<{ id: string }>(sql`SELECT id FROM workspace_file_search_build From 5a000059dede75d2f564ff04bb2a0f1b97a9d468 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 00:04:42 -0700 Subject: [PATCH 2/2] fix(file-search): measure the cleanup budget after the connection is held Hoisting the remaining-budget read out of the transaction callback made it describe the moment the batch was admitted rather than the moment its statements begin. Time spent waiting for a pooled connection then went unaccounted, and the batch installed a timeout larger than the budget actually left. Keep the cheap check before opening a transaction, and re-read once the connection is in hand so the installed timeout is the budget that remains. --- .../search/chunks.integration.ts | 23 +++++++++++++++++++ .../lib/workspace-files/search/index-state.ts | 6 +++-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/apps/sim/lib/workspace-files/search/chunks.integration.ts b/apps/sim/lib/workspace-files/search/chunks.integration.ts index c957b18e8a2..80969864f3e 100644 --- a/apps/sim/lib/workspace-files/search/chunks.integration.ts +++ b/apps/sim/lib/workspace-files/search/chunks.integration.ts @@ -396,6 +396,29 @@ describe('chunked workspace file search on PostgreSQL', () => { (await connection`SELECT count(*)::int AS count FROM workspace_file_search_chunk`)[0].count ).toBe(1000) }) + it('abandons a batch whose budget was spent acquiring its connection', async () => { + const build = (await beginFileSearchBuild(revision))! + await connection`INSERT INTO workspace_file_search_chunk (build_id, workspace_id, ordinal, line_start, fragment, content) + SELECT ${build.id}, 'workspace-1', n, n + 1, false, 'x' FROM generate_series(0, 999) n` + await connection`UPDATE workspace_file_search_build SET expires_at = now() WHERE id = ${build.id}` + + /** Full budget when the batch is admitted, none left once its connection is in hand. */ + const startedAt = Date.now() + const clock = vi + .spyOn(Date, 'now') + .mockReturnValueOnce(startedAt) + .mockReturnValueOnce(startedAt) + .mockReturnValue(startedAt + FILE_SEARCH_CLEANUP_BUDGET_MS - 1) + try { + await expect(cleanupFileSearchBuilds()).resolves.toBe(0) + } finally { + clock.mockRestore() + } + + expect( + (await connection`SELECT count(*)::int AS count FROM workspace_file_search_chunk`)[0].count + ).toBe(1000) + }) it('retires many small builds within one cleanup run', async () => { await connection`INSERT INTO workspace_file_search_build (id, file_id, workspace_id, source_content_updated_at, expires_at) SELECT 'retired-' || n, 'file-1', 'workspace-1', now(), now() FROM generate_series(1, 100) n` diff --git a/apps/sim/lib/workspace-files/search/index-state.ts b/apps/sim/lib/workspace-files/search/index-state.ts index 98a748efa72..9e04446f495 100644 --- a/apps/sim/lib/workspace-files/search/index-state.ts +++ b/apps/sim/lib/workspace-files/search/index-state.ts @@ -259,9 +259,11 @@ export async function cleanupFileSearchBuilds(): Promise { const deadline = Date.now() + FILE_SEARCH_CLEANUP_BUDGET_MS let deleted = 0 for (let batch = 0; batch < FILE_SEARCH_CLEANUP_MAX_BATCHES; batch++) { - const remainingBudget = deadline - Date.now() - if (remainingBudget < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) break + if (deadline - Date.now() < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) break const result = await db.transaction(async (tx) => { + /** Re-read: acquiring the connection can itself have spent the rest of the budget. */ + const remainingBudget = deadline - Date.now() + if (remainingBudget < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) return null await configureFileSearchTransaction(tx, { statementTimeout: remainingBudget }) const builds = await tx.execute<{ id: string