Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 70 additions & 0 deletions apps/sim/lib/workspace-files/search/dispatcher.integration.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 = () => {}
Expand Down
72 changes: 30 additions & 42 deletions apps/sim/lib/workspace-files/search/dispatcher.ts
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,7 @@ async function reapStaleClaims(tx: DbTransaction, now: Date): Promise<number> {
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[],
Expand All @@ -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
Expand Down
Loading