Skip to content

Commit 8508ba1

Browse files
icecrasher321icecrasher321
authored
fix(file-search): bound dispatcher claims and skip locked rows (#7940)
1 parent 1a7d66f commit 8508ba1

2 files changed

Lines changed: 100 additions & 42 deletions

File tree

‎apps/sim/lib/workspace-files/search/dispatcher.integration.ts‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,11 @@ describe('workspace file search dispatch PostgreSQL deadlines', () => {
6161
workspace_id text PRIMARY KEY, enqueued_at timestamp NOT NULL,
6262
updated_at timestamp NOT NULL, last_dispatched_at timestamp
6363
)`
64+
await connection`CREATE INDEX ON workspace_file_search_index
65+
(workspace_id, updated_at, file_id, source_content_updated_at)
66+
WHERE status = 'pending' AND dispatched_at IS NULL`
67+
await connection`CREATE INDEX ON workspace_file_search_index (workspace_id, dispatched_at)
68+
WHERE status = 'pending' AND dispatched_at IS NOT NULL`
6469
await connection`INSERT INTO workspace_file_search_backfill (id, updated_at)
6570
VALUES ('workspace-file-search-v1', '2026-09-16 00:00:00')`
6671
database.current = drizzle(connection)
@@ -92,6 +97,71 @@ describe('workspace file search dispatch PostgreSQL deadlines', () => {
9297
})
9398
}
9499

100+
async function seedQueue(workspaceId: string, queued: number, active = 0) {
101+
await connection`UPDATE workspace_file_search_backfill SET completed_at = now()`
102+
await connection`INSERT INTO workspace_files (id, workspace_id, context, content_updated_at)
103+
SELECT ${workspaceId} || '-' || lpad(n::text, 6, '0'), ${workspaceId}, 'workspace', '2026-09-16'
104+
FROM generate_series(1, ${queued + active}) n`
105+
await connection`INSERT INTO workspace_file_search_index
106+
(file_id, workspace_id, source_content_updated_at, status, updated_at, dispatched_at)
107+
SELECT id, workspace_id, content_updated_at, 'pending', '2026-09-16',
108+
CASE WHEN row_number() OVER (ORDER BY id DESC) <= ${active} THEN now() ELSE NULL END
109+
FROM workspace_files WHERE workspace_id = ${workspaceId}`
110+
await connection`INSERT INTO workspace_file_search_dispatch_queue
111+
(workspace_id, enqueued_at, updated_at) VALUES (${workspaceId}, now(), now())`
112+
}
113+
114+
it('skips a locked candidate without losing it or exceeding workspace capacity', async () => {
115+
await seedQueue('workspace-1', 3, 1)
116+
await connection.begin(async (tx) => {
117+
await tx`SELECT file_id FROM workspace_file_search_index
118+
WHERE file_id = 'workspace-1-000001' FOR UPDATE`
119+
const results = await Promise.all([
120+
prepareWorkspaceFileSearchDispatch(),
121+
prepareWorkspaceFileSearchDispatch(),
122+
])
123+
expect(results.flatMap((result) => result.payloads).map((payload) => payload.fileId)).toEqual(
124+
['workspace-1-000002']
125+
)
126+
const [locked] = await tx`SELECT dispatched_at FROM workspace_file_search_index
127+
WHERE file_id = 'workspace-1-000001'`
128+
expect(locked.dispatched_at).toBeNull()
129+
})
130+
expect((await prepareWorkspaceFileSearchDispatch()).payloads).toEqual([])
131+
await connection`UPDATE workspace_file_search_index SET status = 'ready'
132+
WHERE file_id = 'workspace-1-000002'`
133+
const retry = await prepareWorkspaceFileSearchDispatch()
134+
expect(retry.payloads.map((payload) => payload.fileId)).toEqual(['workspace-1-000001'])
135+
})
136+
137+
it('claims only available slots from a large backlog and preserves current-file eligibility', async () => {
138+
await seedQueue('workspace-1', 10_000, 1)
139+
await seedQueue('workspace-2', 3)
140+
await connection`UPDATE workspace_files SET deleted_at = now() WHERE id = 'workspace-1-000001'`
141+
await connection`UPDATE workspace_files SET content_updated_at = '2026-09-17'
142+
WHERE id = 'workspace-1-000002'`
143+
await connection`UPDATE workspace_files SET context = 'execution' WHERE id = 'workspace-1-000003'`
144+
const result = await prepareWorkspaceFileSearchDispatch()
145+
expect(result.payloads.map((payload) => payload.fileId).sort()).toEqual([
146+
'workspace-1-000004',
147+
'workspace-2-000001',
148+
'workspace-2-000002',
149+
])
150+
expect((await prepareWorkspaceFileSearchDispatch()).payloads).toEqual([])
151+
})
152+
153+
it('honors the remaining global capacity across workspace probes', async () => {
154+
await seedQueue('workspace-1', 3)
155+
await seedQueue('workspace-2', 3)
156+
await seedQueue('workspace-active', 0, 99)
157+
const result = await prepareWorkspaceFileSearchDispatch()
158+
expect(result.payloads).toHaveLength(1)
159+
expect((await prepareWorkspaceFileSearchDispatch()).payloads).toEqual([])
160+
const [row] = await connection`SELECT count(*)::int AS active FROM workspace_file_search_index
161+
WHERE status = 'pending' AND dispatched_at IS NOT NULL`
162+
expect(row.active).toBe(100)
163+
})
164+
95165
it('fails on a locked backfill row and releases the dispatcher lock', async () => {
96166
let release = () => {}
97167
let locked = () => {}

‎apps/sim/lib/workspace-files/search/dispatcher.ts‎

Lines changed: 30 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,7 @@ async function reapStaleClaims(tx: DbTransaction, now: Date): Promise<number> {
301301
return rows.length
302302
}
303303

304+
/** Probe each workspace's available slots and lock candidates before the update, skipping busy rows. */
304305
async function claimQueuedWorkspaceJobs(
305306
tx: DbTransaction,
306307
workspaceIds: readonly string[],
@@ -320,48 +321,35 @@ async function claimQueuedWorkspaceJobs(
320321
WITH selected_workspace(workspace_id) AS (
321322
VALUES ${workspaceValues}
322323
),
323-
workspace_active AS (
324-
SELECT search_index.workspace_id, count(*)::int AS active_count
325-
FROM workspace_file_search_index AS search_index
326-
INNER JOIN selected_workspace AS selected
327-
ON selected.workspace_id = search_index.workspace_id
328-
WHERE search_index.status = 'pending'
329-
AND search_index.dispatched_at IS NOT NULL
330-
GROUP BY search_index.workspace_id
331-
),
332-
ranked AS (
333-
SELECT
334-
search_index.workspace_id,
335-
search_index.file_id,
336-
search_index.source_content_updated_at,
337-
search_index.updated_at,
338-
coalesce(workspace_active.active_count, 0) AS active_count,
339-
row_number() OVER (
340-
PARTITION BY search_index.workspace_id
341-
ORDER BY
342-
search_index.updated_at,
343-
search_index.file_id,
344-
search_index.source_content_updated_at
345-
) AS workspace_rank
346-
FROM workspace_file_search_index AS search_index
347-
INNER JOIN selected_workspace AS selected
348-
ON selected.workspace_id = search_index.workspace_id
349-
INNER JOIN workspace_files AS file
350-
ON file.id = search_index.file_id
351-
AND file.workspace_id = search_index.workspace_id
352-
AND file.context = 'workspace'
353-
AND file.deleted_at IS NULL
354-
AND file.content_updated_at = search_index.source_content_updated_at
355-
LEFT JOIN workspace_active
356-
ON workspace_active.workspace_id = search_index.workspace_id
357-
WHERE search_index.status = 'pending'
358-
AND search_index.dispatched_at IS NULL
359-
),
360-
candidates AS (
361-
SELECT workspace_id, file_id, source_content_updated_at
362-
FROM ranked
363-
WHERE workspace_rank <= ${FILE_SEARCH_INDEX_WORKSPACE_OUTSTANDING} - active_count
364-
ORDER BY updated_at, workspace_id, file_id, source_content_updated_at
324+
candidates AS MATERIALIZED (
325+
SELECT queued.*
326+
FROM selected_workspace AS selected
327+
CROSS JOIN LATERAL (
328+
SELECT count(*)::int AS active_count
329+
FROM (
330+
SELECT 1 FROM workspace_file_search_index AS active
331+
WHERE active.workspace_id = selected.workspace_id
332+
AND active.status = 'pending' AND active.dispatched_at IS NOT NULL
333+
LIMIT ${FILE_SEARCH_INDEX_WORKSPACE_OUTSTANDING}
334+
) AS active_claims
335+
) AS workspace_active
336+
CROSS JOIN LATERAL (
337+
SELECT search_index.workspace_id, search_index.file_id,
338+
search_index.source_content_updated_at, search_index.updated_at
339+
FROM workspace_file_search_index AS search_index
340+
INNER JOIN workspace_files AS file
341+
ON file.id = search_index.file_id
342+
AND file.workspace_id = search_index.workspace_id
343+
AND file.context = 'workspace'
344+
AND file.deleted_at IS NULL
345+
AND file.content_updated_at = search_index.source_content_updated_at
346+
WHERE search_index.workspace_id = selected.workspace_id
347+
AND search_index.status = 'pending' AND search_index.dispatched_at IS NULL
348+
ORDER BY search_index.updated_at, search_index.file_id, search_index.source_content_updated_at
349+
LIMIT greatest(0, ${FILE_SEARCH_INDEX_WORKSPACE_OUTSTANDING} - workspace_active.active_count)
350+
FOR UPDATE OF search_index SKIP LOCKED
351+
) AS queued
352+
ORDER BY queued.updated_at, queued.workspace_id, queued.file_id, queued.source_content_updated_at
365353
LIMIT ${remainingGlobalCapacity}
366354
)
367355
UPDATE workspace_file_search_index AS search_index

0 commit comments

Comments
 (0)