diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index ee25a9bc578..97b11d4917d 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -134,6 +134,13 @@ jobs: BILLING_USAGE_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5433/sim_billing_test run: bunx vitest run lib/billing/core/usage-log.postgres.test.ts + - name: Verify file search dispatch deadlines on PostgreSQL 17 + working-directory: apps/sim + env: + TZ: America/Los_Angeles + KNOWLEDGE_ACL_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_auth_scim + run: bunx vitest run --mode integration lib/workspace-files/search/dispatcher.integration.ts + - name: Verify SCIM and administration over real HTTP working-directory: apps/sim env: diff --git a/apps/sim/lib/workspace-files/search/constants.ts b/apps/sim/lib/workspace-files/search/constants.ts index 20b6047c39e..f4a2820e663 100644 --- a/apps/sim/lib/workspace-files/search/constants.ts +++ b/apps/sim/lib/workspace-files/search/constants.ts @@ -42,6 +42,10 @@ export const FILE_SEARCH_INDEX_MAX_OUTSTANDING = 100 export const FILE_SEARCH_INDEX_DISPATCH_WORKSPACES = 100 export const FILE_SEARCH_DISPATCH_INTERVAL_MS = 60 * 1000 export const FILE_SEARCH_DISPATCH_MAX_DURATION_SECONDS = 60 +/** Leave room for connection setup, rollback, and task failure reporting before the hard cutoff. */ +export const FILE_SEARCH_DISPATCH_STATEMENT_TIMEOUT_MS = 10 * 1000 +export const FILE_SEARCH_DISPATCH_LOCK_TIMEOUT_MS = 2 * 1000 +export const FILE_SEARCH_DISPATCH_TRANSACTION_TIMEOUT_MS = 20 * 1000 export const FILE_SEARCH_INDEX_MAX_DURATION_SECONDS = 15 * 60 export const FILE_SEARCH_INDEX_STALE_DISPATCH_MS = 6 * 60 * 60 * 1000 export const FILE_SEARCH_INDEX_STALE_REAP_LIMIT = 100 diff --git a/apps/sim/lib/workspace-files/search/dispatcher.integration.ts b/apps/sim/lib/workspace-files/search/dispatcher.integration.ts new file mode 100644 index 00000000000..2a9eadf2599 --- /dev/null +++ b/apps/sim/lib/workspace-files/search/dispatcher.integration.ts @@ -0,0 +1,193 @@ +/** Real PostgreSQL cancellation must roll back preparation and release its advisory lock. */ +import { withUtcTimestamps } from '@sim/db/timestamps' +import { getPostgresErrorCode } from '@sim/utils/errors' +import { generateId } from '@sim/utils/id' +import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js' +import postgres from 'postgres' +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const database = vi.hoisted(() => ({ current: undefined as PostgresJsDatabase | undefined })) +const mocks = vi.hoisted(() => ({ batchTrigger: vi.fn() })) +vi.mock('@sim/db', () => ({ + get db() { + if (!database.current) throw new Error('Dispatcher test database is not initialized') + return database.current + }, +})) +vi.mock('@/lib/workspace-files/search/indexing', () => ({ + indexWorkspaceFileForSearch: vi.fn(), + markWorkspaceFileSearchIndexFailed: vi.fn(), +})) +vi.mock('@trigger.dev/sdk', () => ({ tasks: { batchTrigger: mocks.batchTrigger } })) +vi.mock('@/lib/core/config/env-flags', () => ({ isTriggerDevEnabled: true })) +vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' })) + +import { + dispatchWorkspaceFileSearchIndexJobs, + prepareWorkspaceFileSearchDispatch, +} from '@/lib/workspace-files/search/dispatcher' + +describe('workspace file search dispatch PostgreSQL deadlines', () => { + const schemaName = `dispatch_test_${generateId().replaceAll('-', '')}` + const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL + if (!databaseUrl) throw new Error('Dispatcher tests require a disposable local database') + const connection = postgres( + databaseUrl, + withUtcTimestamps({ + max: 3, + prepare: false, + fetch_types: false, + connection: { search_path: schemaName }, + onnotice: () => {}, + }) + ) + + beforeAll(async () => { + await connection`CREATE SCHEMA ${connection(schemaName)}` + await connection`CREATE TABLE workspace_file_search_backfill ( + id text PRIMARY KEY, after_workspace_id text, after_file_id text, + completed_at timestamp, updated_at timestamp NOT NULL + )` + await connection`CREATE TABLE workspace_files ( + id text PRIMARY KEY, workspace_id text NOT NULL, context text NOT NULL, + deleted_at timestamp, content_updated_at timestamp NOT NULL + )` + await connection`CREATE TABLE workspace_file_search_index ( + file_id text NOT NULL, workspace_id text NOT NULL, source_content_updated_at timestamp NOT NULL, + status text NOT NULL, dispatched_at timestamp, updated_at timestamp NOT NULL, + PRIMARY KEY (file_id, source_content_updated_at) + )` + await connection`CREATE TABLE workspace_file_search_dispatch_queue ( + workspace_id text PRIMARY KEY, enqueued_at timestamp NOT NULL, + updated_at timestamp NOT NULL, last_dispatched_at timestamp + )` + 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) + }) + + beforeEach(async () => { + mocks.batchTrigger.mockReset() + await connection`DROP TRIGGER IF EXISTS slow_backfill ON workspace_file_search_backfill` + await connection`TRUNCATE workspace_files, workspace_file_search_index, workspace_file_search_dispatch_queue` + await connection`UPDATE workspace_file_search_backfill + SET updated_at = '2026-09-16 00:00:00', completed_at = NULL` + }) + + afterAll(async () => { + try { + await connection`DROP SCHEMA ${connection(schemaName)} CASCADE` + } finally { + await connection.end() + database.current = undefined + } + }) + + async function expectAdvisoryLockReleased() { + await connection.begin(async (tx) => { + const [row] = await tx`SELECT pg_try_advisory_xact_lock( + hashtextextended('workspace-file-search-dispatch', 0) + ) AS acquired` + expect(row.acquired).toBe(true) + }) + } + + it('fails on a locked backfill row and releases the dispatcher lock', async () => { + let release = () => {} + let locked = () => {} + const releaseLock = new Promise((resolve) => { + release = resolve + }) + const lockReady = new Promise((resolve) => { + locked = resolve + }) + const blocker = connection.begin(async (tx) => { + await tx`SELECT id FROM workspace_file_search_backfill FOR UPDATE` + locked() + await releaseLock + }) + await lockReady + try { + const failure = await prepareWorkspaceFileSearchDispatch().catch((error: unknown) => error) + expect(getPostgresErrorCode(failure)).toBe('55P03') + await expectAdvisoryLockReleased() + } finally { + release() + await blocker + } + }) + + it('cancels a slow statement and rolls back its earlier writes', async () => { + await connection`CREATE FUNCTION slow_backfill() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN + UPDATE workspace_file_search_backfill SET updated_at = '2099-01-01'; + PERFORM pg_sleep(15); + RETURN NEW; + END + $$` + await connection`CREATE TRIGGER slow_backfill BEFORE INSERT ON workspace_file_search_backfill + FOR EACH ROW EXECUTE FUNCTION slow_backfill()` + + const failure = await prepareWorkspaceFileSearchDispatch().catch((error: unknown) => error) + + expect(getPostgresErrorCode(failure)).toBe('57014') + const [row] = + await connection`SELECT updated_at::text AS updated_at FROM workspace_file_search_backfill` + expect(row.updated_at).toBe('2026-09-16 00:00:00') + await expectAdvisoryLockReleased() + }, 20_000) + + it('releases committed claims without the preparation deadlines', async () => { + const fileId = generateId() + const workspaceId = generateId() + await connection`UPDATE workspace_file_search_backfill SET completed_at = now()` + await connection`INSERT INTO workspace_files (id, workspace_id, context, content_updated_at) + VALUES (${fileId}, ${workspaceId}, 'workspace', '2026-09-16')` + await connection`INSERT INTO workspace_file_search_index + (file_id, workspace_id, source_content_updated_at, status, updated_at) + VALUES (${fileId}, ${workspaceId}, '2026-09-16', 'pending', now())` + await connection`INSERT INTO workspace_file_search_dispatch_queue + (workspace_id, enqueued_at, updated_at) VALUES (${workspaceId}, now(), now())` + await connection`CREATE TABLE cleanup_timeouts ( + lock_timeout text, statement_timeout text, transaction_timeout text + )` + await connection`CREATE FUNCTION record_cleanup_timeouts() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN + INSERT INTO cleanup_timeouts VALUES ( + current_setting('lock_timeout'), + current_setting('statement_timeout'), + current_setting('transaction_timeout') + ); + RETURN NEW; + END + $$` + await connection`CREATE TRIGGER record_cleanup_timeouts AFTER UPDATE OF dispatched_at + ON workspace_file_search_index FOR EACH ROW + WHEN (OLD.dispatched_at IS NOT NULL AND NEW.dispatched_at IS NULL) + EXECUTE FUNCTION record_cleanup_timeouts()` + + const enqueueError = new Error('Queue unavailable') + mocks.batchTrigger.mockRejectedValueOnce(enqueueError) + + await expect(dispatchWorkspaceFileSearchIndexJobs()).rejects.toBe(enqueueError) + expect(mocks.batchTrigger).toHaveBeenCalledWith('workspace-file-search-index', [ + expect.objectContaining({ + payload: { + fileId, + workspaceId, + sourceContentUpdatedAt: '2026-09-16T00:00:00.000Z', + }, + }), + ]) + const [index] = await connection`SELECT dispatched_at FROM workspace_file_search_index + WHERE file_id = ${fileId}` + expect(index.dispatched_at).toBeNull() + const [queued] = await connection`SELECT workspace_id FROM workspace_file_search_dispatch_queue + WHERE workspace_id = ${workspaceId}` + expect(queued.workspace_id).toBe(workspaceId) + const timeouts = await connection`SELECT * FROM cleanup_timeouts` + expect([...timeouts]).toEqual([ + { lock_timeout: '0', statement_timeout: '0', transaction_timeout: '0' }, + ]) + }) +}) diff --git a/apps/sim/lib/workspace-files/search/dispatcher.test.ts b/apps/sim/lib/workspace-files/search/dispatcher.test.ts index 98845cb5eff..937c30f4277 100644 --- a/apps/sim/lib/workspace-files/search/dispatcher.test.ts +++ b/apps/sim/lib/workspace-files/search/dispatcher.test.ts @@ -1,9 +1,43 @@ /** * @vitest-environment node */ -import { describe, expect, it } from 'vitest' +import { + workspaceFileSearchBackfill, + workspaceFileSearchDispatchQueue, + workspaceFileSearchIndex, +} from '@sim/db/schema' +import { dbChainMock, dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + batchTrigger: vi.fn(), + info: vi.fn(), + error: vi.fn(), +})) + +vi.mock('@sim/db/schema', async () => ({ + ...(await import('@sim/testing/mocks/schema.mock')).schemaMock, + workspaceFileSearchBackfill: { id: 'backfill.id' }, + workspaceFileSearchDispatchQueue: { + workspaceId: 'queue.workspaceId', + lastDispatchedAt: 'queue.lastDispatchedAt', + enqueuedAt: 'queue.enqueuedAt', + }, +})) + +vi.mock('@sim/logger', () => ({ createLogger: () => ({ info: mocks.info, error: mocks.error }) })) +vi.mock('@trigger.dev/sdk', () => ({ tasks: { batchTrigger: mocks.batchTrigger } })) +vi.mock('@/lib/core/config/env-flags', () => ({ isTriggerDevEnabled: true })) +vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' })) +vi.mock('@/lib/workspace-files/search/indexing', () => ({ + indexWorkspaceFileForSearch: vi.fn(), + markWorkspaceFileSearchIndexFailed: vi.fn(), +})) + import { buildWorkspaceFileSearchTriggerItems, + dispatchWorkspaceFileSearchIndexJobs, + prepareWorkspaceFileSearchDispatch, shouldUseWorkspaceFileSearchTrigger, } from '@/lib/workspace-files/search/dispatcher' @@ -34,3 +68,119 @@ describe('workspace file search dispatch policy', () => { ]) }) }) + +describe('workspace file search dispatch deadlines', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + + it('sets local database deadlines before taking the advisory lock', async () => { + dbChainMockFns.execute.mockResolvedValueOnce([]).mockResolvedValueOnce([{ acquired: false }]) + + await expect(prepareWorkspaceFileSearchDispatch()).resolves.toEqual({ + payloads: [], + backfilledFiles: 0, + reapedClaims: 0, + lockAcquired: false, + }) + + const guards = JSON.stringify(dbChainMockFns.execute.mock.calls[0][0]) + expect(guards).toContain("set_config('statement_timeout', ") + expect(guards).toContain('10000ms') + expect(guards).toContain("set_config('lock_timeout', ") + expect(guards).toContain('2000ms') + expect(guards).toContain("'transaction_timeout'") + expect(guards).toContain('20000ms') + expect(JSON.stringify(dbChainMockFns.execute.mock.calls[1][0])).toContain( + 'pg_try_advisory_xact_lock' + ) + expect(dbChainMockFns.insert).not.toHaveBeenCalled() + }) + + it.each(['57014', '55P03', '25P04'])( + 'propagates SQLSTATE %s without enqueuing an uncommitted claim', + async (code) => { + const error = new Error('Failed query\nparams: sensitive-value', { + cause: Object.assign(new Error('database timeout'), { code }), + }) + dbChainMockFns.execute.mockResolvedValueOnce([]).mockResolvedValueOnce([{ acquired: true }]) + dbChainMockFns.onConflictDoNothing.mockRejectedValueOnce(error) + + await expect(dispatchWorkspaceFileSearchIndexJobs()).rejects.toBe(error) + + expect(mocks.batchTrigger).not.toHaveBeenCalled() + expect(mocks.error).toHaveBeenCalledWith('Workspace file search dispatch phase failed', { + phase: 'backfill', + durationMs: expect.any(Number), + code, + error: 'Failed query', + }) + expect(JSON.stringify(mocks.error.mock.calls)).not.toContain('sensitive-value') + } + ) + + it('reports a transaction failure even after the transaction callback finishes', async () => { + const error = Object.assign(new Error('commit failed'), { code: '08006' }) + dbChainMockFns.execute.mockResolvedValueOnce([]).mockResolvedValueOnce([{ acquired: false }]) + dbChainMockFns.transaction.mockImplementationOnce(async (callback) => { + await callback(dbChainMock.db) + throw error + }) + + await expect(prepareWorkspaceFileSearchDispatch()).rejects.toBe(error) + + expect(mocks.error).toHaveBeenCalledWith('Workspace file search dispatch phase failed', { + phase: 'prepare-transaction', + durationMs: expect.any(Number), + code: '08006', + error: 'commit failed', + }) + }) + + it.each([false, true])( + 'preserves enqueue failures when claim release fails: %s', + async (releaseFails) => { + queueTableRows(workspaceFileSearchBackfill, [{ completedAt: new Date() }]) + queueTableRows(workspaceFileSearchIndex, []) + queueTableRows(workspaceFileSearchIndex, [{ active: 0 }]) + queueTableRows(workspaceFileSearchDispatchQueue, [{ workspaceId: 'workspace-1' }]) + dbChainMockFns.execute + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([{ acquired: true }]) + .mockResolvedValueOnce([ + { + workspaceId: 'workspace-1', + fileId: 'file-1', + sourceContentUpdatedAt: new Date('2026-09-16T00:00:00Z'), + }, + ]) + const error = new Error('Trigger unavailable') + const releaseError = new Error('claim release unavailable') + mocks.batchTrigger.mockRejectedValueOnce(error) + if (releaseFails) { + dbChainMockFns.transaction + .mockImplementationOnce(async (callback) => callback(dbChainMock.db)) + .mockRejectedValueOnce(releaseError) + } + + if (releaseFails) { + await expect(dispatchWorkspaceFileSearchIndexJobs()).rejects.toMatchObject({ + errors: [error, releaseError], + cause: error, + }) + } else { + await expect(dispatchWorkspaceFileSearchIndexJobs()).rejects.toBe(error) + expect(dbChainMockFns.set).toHaveBeenCalledWith( + expect.objectContaining({ dispatchedAt: null }) + ) + } + + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2) + const guards = dbChainMockFns.execute.mock.calls.filter(([query]) => + JSON.stringify(query).includes('statement_timeout') + ) + expect(guards).toHaveLength(1) + } + ) +}) diff --git a/apps/sim/lib/workspace-files/search/dispatcher.ts b/apps/sim/lib/workspace-files/search/dispatcher.ts index 3c9518a9fe1..91f4b6d1899 100644 --- a/apps/sim/lib/workspace-files/search/dispatcher.ts +++ b/apps/sim/lib/workspace-files/search/dispatcher.ts @@ -7,7 +7,8 @@ import { workspaceFiles, } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { getErrorMessage } from '@sim/utils/errors' +import { getErrorMessage, getPostgresErrorCode } from '@sim/utils/errors' +import { truncate } from '@sim/utils/string' import { and, asc, @@ -30,6 +31,9 @@ import { runDetached } from '@/lib/core/utils/background' import type { DbTransaction } from '@/lib/db/types' import { FILE_SEARCH_BACKFILL_PAGE_SIZE, + FILE_SEARCH_DISPATCH_LOCK_TIMEOUT_MS, + FILE_SEARCH_DISPATCH_STATEMENT_TIMEOUT_MS, + FILE_SEARCH_DISPATCH_TRANSACTION_TIMEOUT_MS, FILE_SEARCH_INDEX_DISPATCH_WORKSPACES, FILE_SEARCH_INDEX_MAX_OUTSTANDING, FILE_SEARCH_INDEX_STALE_DISPATCH_MS, @@ -47,6 +51,41 @@ const logger = createLogger('WorkspaceFileSearchDispatcher') const DISPATCH_LOCK_NAME = 'workspace-file-search-dispatch' const BACKFILL_CURSOR_ID = 'workspace-file-search-v1' +async function runDispatchPhase(phase: string, operation: () => Promise): Promise { + const startedAt = Date.now() + logger.info('Workspace file search dispatch phase started', { phase }) + try { + const result = await operation() + logger.info('Workspace file search dispatch phase completed', { + phase, + durationMs: Date.now() - startedAt, + }) + return result + } catch (error) { + logger.error('Workspace file search dispatch phase failed', { + phase, + durationMs: Date.now() - startedAt, + code: getPostgresErrorCode(error), + error: truncate(getErrorMessage(error).split('\nparams: ')[0], 500), + }) + throw error + } +} + +/** Preparation must roll back before the worker's hard deadline. */ +async function configureDispatchTimeouts(tx: DbTransaction): Promise { + await tx.execute(sql` + SELECT + set_config('statement_timeout', ${`${FILE_SEARCH_DISPATCH_STATEMENT_TIMEOUT_MS}ms`}, true), + set_config('lock_timeout', ${`${FILE_SEARCH_DISPATCH_LOCK_TIMEOUT_MS}ms`}, true), + set_config( + 'transaction_timeout', + ${`${FILE_SEARCH_DISPATCH_TRANSACTION_TIMEOUT_MS}ms`}, + true + ) + `) +} + interface RevisionIdentity { fileId: string sourceContentUpdatedAt: Date @@ -276,7 +315,7 @@ async function claimQueuedWorkspaceJobs( const rows = await tx.execute<{ workspaceId: string fileId: string - sourceContentUpdatedAt: Date + sourceContentUpdatedAt: string }>(sql` WITH selected_workspace(workspace_id) AS ( VALUES ${workspaceValues} @@ -335,7 +374,7 @@ async function claimQueuedWorkspaceJobs( RETURNING search_index.workspace_id AS "workspaceId", search_index.file_id AS "fileId", - search_index.source_content_updated_at AS "sourceContentUpdatedAt" + search_index.source_content_updated_at AT TIME ZONE 'UTC' AS "sourceContentUpdatedAt" `) const remainingForWorkspace = tx @@ -384,50 +423,60 @@ async function claimQueuedWorkspaceJobs( } export async function prepareWorkspaceFileSearchDispatch(): Promise { - return db.transaction(async (tx) => { - const [lock] = await tx.execute<{ acquired: boolean }>( - sql`SELECT pg_try_advisory_xact_lock(hashtextextended(${DISPATCH_LOCK_NAME}, 0)) AS acquired` - ) - if (!lock?.acquired) { - return { payloads: [], backfilledFiles: 0, reapedClaims: 0, lockAcquired: false } - } + return runDispatchPhase('prepare-transaction', () => + db.transaction(async (tx) => { + await runDispatchPhase('configure-timeouts', () => configureDispatchTimeouts(tx)) + return runDispatchPhase('prepare', async () => { + const [lock] = await tx.execute<{ acquired: boolean }>( + sql`SELECT pg_try_advisory_xact_lock(hashtextextended(${DISPATCH_LOCK_NAME}, 0)) AS acquired` + ) + if (!lock?.acquired) { + return { payloads: [], backfilledFiles: 0, reapedClaims: 0, lockAcquired: false } + } - const now = new Date() - const backfilledFiles = await seedBackfillPage(tx, now) - const reapedClaims = await reapStaleClaims(tx, now) - const [{ active }] = await tx - .select({ active: count() }) - .from(workspaceFileSearchIndex) - .where( - and( - eq(workspaceFileSearchIndex.status, 'pending'), - isNotNull(workspaceFileSearchIndex.dispatchedAt) + const now = new Date() + const backfilledFiles = await runDispatchPhase('backfill', () => seedBackfillPage(tx, now)) + const reapedClaims = await runDispatchPhase('reap', () => reapStaleClaims(tx, now)) + const [{ active }] = await tx + .select({ active: count() }) + .from(workspaceFileSearchIndex) + .where( + and( + eq(workspaceFileSearchIndex.status, 'pending'), + isNotNull(workspaceFileSearchIndex.dispatchedAt) + ) + ) + const remainingGlobalCapacity = Math.max( + 0, + FILE_SEARCH_INDEX_MAX_OUTSTANDING - Number(active) ) - ) - const remainingGlobalCapacity = Math.max(0, FILE_SEARCH_INDEX_MAX_OUTSTANDING - Number(active)) - if (remainingGlobalCapacity === 0) { - return { payloads: [], backfilledFiles, reapedClaims, lockAcquired: true } - } + if (remainingGlobalCapacity === 0) { + return { payloads: [], backfilledFiles, reapedClaims, lockAcquired: true } + } - const workspaces = await tx - .select({ workspaceId: workspaceFileSearchDispatchQueue.workspaceId }) - .from(workspaceFileSearchDispatchQueue) - .orderBy( - sql`${workspaceFileSearchDispatchQueue.lastDispatchedAt} ASC NULLS FIRST`, - asc(workspaceFileSearchDispatchQueue.enqueuedAt), - asc(workspaceFileSearchDispatchQueue.workspaceId) - ) - .limit(Math.min(FILE_SEARCH_INDEX_DISPATCH_WORKSPACES, remainingGlobalCapacity)) - .for('update', { skipLocked: true }) + const workspaces = await tx + .select({ workspaceId: workspaceFileSearchDispatchQueue.workspaceId }) + .from(workspaceFileSearchDispatchQueue) + .orderBy( + sql`${workspaceFileSearchDispatchQueue.lastDispatchedAt} ASC NULLS FIRST`, + asc(workspaceFileSearchDispatchQueue.enqueuedAt), + asc(workspaceFileSearchDispatchQueue.workspaceId) + ) + .limit(Math.min(FILE_SEARCH_INDEX_DISPATCH_WORKSPACES, remainingGlobalCapacity)) + .for('update', { skipLocked: true }) - const payloads = await claimQueuedWorkspaceJobs( - tx, - workspaces.map((workspace) => workspace.workspaceId), - remainingGlobalCapacity, - now - ) - return { payloads, backfilledFiles, reapedClaims, lockAcquired: true } - }) + const payloads = await runDispatchPhase('claim', () => + claimQueuedWorkspaceJobs( + tx, + workspaces.map((workspace) => workspace.workspaceId), + remainingGlobalCapacity, + now + ) + ) + return { payloads, backfilledFiles, reapedClaims, lockAcquired: true } + }) + }) + ) } async function releaseDispatchClaims(payloads: readonly WorkspaceFileSearchIndexPayload[]) { @@ -437,20 +486,22 @@ async function releaseDispatchClaims(payloads: readonly WorkspaceFileSearchIndex fileId: payload.fileId, sourceContentUpdatedAt: new Date(payload.sourceContentUpdatedAt), })) - await db.transaction(async (tx) => { - const filter = revisionFilter(rows) - if (filter) { - await tx - .update(workspaceFileSearchIndex) - .set({ dispatchedAt: null, updatedAt: new Date() }) - .where(and(filter, eq(workspaceFileSearchIndex.status, 'pending'))) - } - await enqueueWorkspaces( - tx, - rows.map((row) => row.workspaceId), - new Date() - ) - }) + await runDispatchPhase('release-claims', () => + db.transaction(async (tx) => { + const filter = revisionFilter(rows) + if (filter) { + await tx + .update(workspaceFileSearchIndex) + .set({ dispatchedAt: null, updatedAt: new Date() }) + .where(and(filter, eq(workspaceFileSearchIndex.status, 'pending'))) + } + await enqueueWorkspaces( + tx, + rows.map((row) => row.workspaceId), + new Date() + ) + }) + ) } async function dispatchPreparedJobs( @@ -497,7 +548,9 @@ export async function dispatchWorkspaceFileSearchIndexJobs(): Promise + dispatchPreparedJobs(prepared.payloads) + ) return { dispatchedFiles, backfilledFiles: prepared.backfilledFiles, @@ -505,11 +558,21 @@ export async function dispatchWorkspaceFileSearchIndexJobs(): Promise