diff --git a/apps/sim/lib/workspace-files/search/chunks.integration.ts b/apps/sim/lib/workspace-files/search/chunks.integration.ts index ee6825ff29c..80969864f3e 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,51 @@ 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('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/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..9e04446f495 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,13 @@ 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++) { + if (deadline - Date.now() < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) break const result = await db.transaction(async (tx) => { - await configureFileSearchTransaction(tx, { - statementTimeout: Math.max(1, deadline - Date.now()), - }) + /** 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 }>(sql`SELECT id FROM workspace_file_search_build