fix(file-search): avoid foreground GIN cleanup stalls

Author: icecrasher321

apps/sim/lib/workspace-files/search/README.md

@@ -10,7 +10,11 @@ PostgreSQL stores complete extracted text in bounded chunks. Object storage rema
 
 8 KiB values may still use PostgreSQL TOAST. The bound controls the size of each logical value and detoast operation; avoiding TOAST entirely is not the objective. Tiny lines share rows, so row count scales with bytes instead of newline count. Worst-case line packing can leave roughly half a block unused; long-line overlap adds at most eight bytes per fragment.
 
-Workers download and extract outside database transactions, then insert batches of at most 250 rows / 1 MiB. Each batch checks the build token and lease. Publication locks the canonical file, build, and revision in that order, verifies the stored chunk count, and changes the visible pointer only after every batch succeeds. Old dispatch failure callbacks cannot overwrite newer dispatches or successful builds.
+Workers download and extract outside database transactions, then insert batches of at most 250 rows / 128 KiB. Each batch checks the build token and lease. Publication locks the canonical file, build, and revision in that order, verifies the stored chunk count, and changes the visible pointer only after every batch succeeds. Old dispatch failure callbacks cannot overwrite newer dispatches or successful builds.
+
+The chunk GIN index uses `fastupdate = off`. Each bounded insert updates the main index directly instead of appending to a shared pending list. With deferred updates enabled, even a small insert can cross the pending-list threshold and synchronously merge accumulated work from other files. Direct updates trade some bulk-write throughput for avoiding that foreground cleanup cliff. They do not eliminate normal index I/O, vacuum, or storage contention; the row and worker limits still apply. The 128 KiB batch budget bounds direct index work per transaction without changing the 25 MiB file coverage limit. Dense text with many distinct trigrams and an index working set larger than the available cache can still exceed the statement deadline. Capacity validation must include that cache pressure, not only a small corpus or a row count.
+
+Indexing transactions have separate limits from search: ten seconds per statement, five seconds waiting for a lock, and thirty seconds total on PostgreSQL 17. The outer limit leaves time for ordinary statement cancellation and rollback instead of terminating the connection at the same ten-second deadline. PostgreSQL 16 uses the compatible idle-transaction guard. A canceled batch remains unpublished; the existing task retry starts a fresh fenced build, and cleanup retires the previous attempt. This does not automatically retry revisions already marked failed.
 
 The indexing task uses an isolated `medium-2x` Trigger worker (4 GB RAM). Document parsers can materialize expanded content before chunking, so source and extracted-text byte limits do not bound parser memory. Parser complexity guards and the worker's memory budget remain separate protections.
 
@@ -41,16 +45,22 @@ Search results retain `fileId`, 1-based `lineNumber`, and bounded `text` preview
 3. Before retiring legacy storage, verify the new app and Trigger workers are fully deployed, old runs/retries have drained, the backfill cursor has completed, and scoped coverage is ready or explicitly excluded. Investigate failed or stale pending revisions. Check cleanup backlog and run representative exact/regex searches, including long lines and folder scopes.
 4. After the rollback window, ship a separate contract PR removing the legacy schema and dropping `workspace_file_search_segment` / `workspace_file_search_index` with a short lock timeout. Do not delete the entire old index row-by-row or backfill it inside the schema migration. Dropping obsolete tables reclaims their heap, indexes, and TOAST together. The `contract-pending` marker in `packages/db/schema.ts` tracks this step.
 
-Until that contract deploy, legacy foreign-key cascades can still make a hard file/workspace deletion expensive. New-index cleanup is bounded; retaining the old schema cannot erase that legacy cost. The earlier timestamp-repair script detects the chunk schema and leaves obsolete legacy text for this contract step instead of deleting it in bulk. No production cleanup is part of this PR.
+Until that contract deploy, legacy foreign-key cascades can still make a hard file/workspace deletion expensive. New-index cleanup is bounded; retaining the old schema cannot erase that legacy cost. The earlier timestamp-repair script detects the chunk schema and leaves obsolete legacy text for this contract step instead of deleting it in bulk. Legacy-table retirement remains a separate contract migration.
 
 Rollback before retirement requires restoring the old trigger function as well as the old app/worker version, and reconciling legacy revisions written during the cutover. Do not assume retained tables are automatically up to date. Canonical revision joins prevent stale content from being returned.
 
+### Direct GIN writes
+
+Migration `0364_workspace_file_search_direct_gin.sql` changes the existing index's storage option without rebuilding it or rewriting chunks. It commits that change before calling PostgreSQL's `gin_clean_pending_list` with a five-minute statement budget. Turning off deferred updates alone leaves the previous pending list in place, so this one-time drain is required. Searches continue to see both existing pending entries and main-index entries throughout the transition, and old workers use the same schema. Deploy the smaller-batch worker policy in the same release; running older workers retain their previous batch size until replaced.
+
+If maintenance times out, the storage option remains off and migration replay safely resumes the drain; no content is discarded and no new pending entries accumulate. The migration uses the runner's direct connection and short DDL lock timeout. It adds no maintenance scheduler. After deployment, check that the index is valid, `reloptions` includes `fastupdate=off`, the migration completed, and chunk insert latency and timeout rates remain healthy under the existing worker concurrency. Previously failed revisions require an explicitly scoped retry after their failure reason and current content version are checked; no blanket backfill or deletion is part of this change.
+
 ## Verification
 
-Run unit tests in `apps/sim` with `bunx vitest run lib/workspace-files/search lib/file-parsers`. Run the PostgreSQL suites on both PostgreSQL 16 and 17 against a disposable local database through `KNOWLEDGE_ACL_TEST_DATABASE_URL` and `--mode integration`. `chunks.integration.ts` applies the actual trigger migrations in an isolated schema. It covers build fencing, revision changes, deletion, cleanup bounds, complete-line matching, UTF-8 boundaries, scope, and admission limits.
+Run unit tests in `apps/sim` with `bunx vitest run lib/workspace-files/search lib/file-parsers`. Run the PostgreSQL suites on both PostgreSQL 16 and 17 against a disposable local database through `KNOWLEDGE_ACL_TEST_DATABASE_URL` and `--mode integration`. `chunks.integration.ts` applies the actual trigger and index migrations in an isolated schema. It covers build fencing, revision changes, deletion, cleanup bounds, complete-line matching, UTF-8 boundaries, scope, admission limits, GIN migration replay with a populated pending list, and statement cancellation followed by a successful retry on an intact connection. Its local database role must be able to install the `pgstattuple` diagnostic extension.
 
 Set `FILE_SEARCH_BENCHMARK_FILES` to change the synthetic file count (default 1,000, maximum 10,000). Set `FILE_SEARCH_BENCHMARK_OUTPUT` to an output path when running the chunk integration suite to record repeated end-to-end searches and `EXPLAIN (ANALYZE, BUFFERS)` plans on a synthetic multi-file corpus. The fixture is synthetic; it contains no production content.
 
 ## PostgreSQL references
 
-The design uses PostgreSQL's documented [TOAST behavior](https://www.postgresql.org/docs/17/storage-toast.html), [trigram index support for LIKE and regex](https://www.postgresql.org/docs/17/pgtrgm.html), and [EXPLAIN guidance](https://www.postgresql.org/docs/17/using-explain.html). The chunk size and query paths are application choices validated by the synthetic fixture, not PostgreSQL hard limits.
+The design uses PostgreSQL's documented [TOAST behavior](https://www.postgresql.org/docs/17/storage-toast.html), [trigram index support for LIKE and regex](https://www.postgresql.org/docs/17/pgtrgm.html), [GIN pending-list tradeoffs](https://www.postgresql.org/docs/17/gin.html#GIN-FAST-UPDATE), [index storage parameters](https://www.postgresql.org/docs/17/sql-createindex.html), and [EXPLAIN guidance](https://www.postgresql.org/docs/17/using-explain.html). The chunk size and query paths are application choices validated by the synthetic fixture, not PostgreSQL hard limits.

apps/sim/lib/workspace-files/search/chunks.integration.ts

@@ -32,9 +32,12 @@ vi.mock('@/lib/copilot/tools/server/files/doc-compile', () => ({ resolveServable
 vi.mock('@/lib/file-parsers', () => ({ parseBuffer: vi.fn(), isSupportedFileType: vi.fn() }))
 
 import {
+  FILE_SEARCH_CHUNK_BYTES,
   FILE_SEARCH_CLEANUP_BATCH_ROWS,
   FILE_SEARCH_CLEANUP_BUDGET_MS,
   FILE_SEARCH_CLEANUP_MAX_BATCHES,
+  FILE_SEARCH_INSERT_BATCH_BYTES,
+  FILE_SEARCH_INSERT_BATCH_ROWS,
   FILE_SEARCH_QUERY_GLOBAL_CONCURRENCY,
   FILE_SEARCH_QUERY_WORKSPACE_CONCURRENCY,
 } from '@/lib/workspace-files/search/constants'
@@ -93,6 +96,35 @@ describe('chunked workspace file search on PostgreSQL', () => {
       onnotice: () => {},
     })
   )
+  const ginWriteMigration = '0364_workspace_file_search_direct_gin.sql'
+
+  async function applyMigration(migration: string) {
+    const source = readFileSync(
+      resolve(process.cwd(), '../../packages/db/migrations', migration),
+      'utf8'
+    ).replaceAll('"public".', `"${schema}".`)
+    const session = await connection.reserve()
+    try {
+      await session`BEGIN`
+      for (const statement of source.split('--> statement-breakpoint'))
+        if (statement.trim()) await session.unsafe(statement)
+      await session`COMMIT`
+    } finally {
+      await session`ROLLBACK`
+      await session`RESET statement_timeout`
+      session.release()
+    }
+  }
+
+  async function ginState() {
+    const [state] =
+      await connection`SELECT c.oid, 'fastupdate=off' = ANY(c.reloptions) AS direct_writes,
+      i.indisvalid, pending.pending_pages
+      FROM pg_class c JOIN pg_index i ON i.indexrelid = c.oid
+      CROSS JOIN LATERAL pgstatginindex(c.oid) pending
+      WHERE c.oid = 'workspace_file_search_chunk_content_idx'::regclass`
+    return state
+  }
 
   async function addFile(
     fileId: string,
@@ -110,10 +142,14 @@ describe('chunked workspace file search on PostgreSQL', () => {
     expect(build).not.toBeNull()
     const plan = planFileSearchIndex({ text, partial: false }, signal)
     const chunks = [...iterateFileSearchChunks(plan, signal)]
-    for (let offset = 0; offset < chunks.length; offset += 100)
-      expect(await appendFileSearchChunks(build!, chunks.slice(offset, offset + 100), signal)).toBe(
-        true
-      )
+    const batchRows = Math.min(
+      FILE_SEARCH_INSERT_BATCH_ROWS,
+      Math.floor(FILE_SEARCH_INSERT_BATCH_BYTES / FILE_SEARCH_CHUNK_BYTES)
+    )
+    for (let offset = 0; offset < chunks.length; offset += batchRows)
+      expect(
+        await appendFileSearchChunks(build!, chunks.slice(offset, offset + batchRows), signal)
+      ).toBe(true)
     expect(
       await publishFileSearchBuild(
         build!,
@@ -140,6 +176,7 @@ describe('chunked workspace file search on PostgreSQL', () => {
   let captureQuery = false
 
   beforeAll(async () => {
+    await connection`CREATE EXTENSION IF NOT EXISTS pgstattuple`
     await connection`CREATE SCHEMA ${connection(schema)}`
     await connection`CREATE TABLE workspace (id text PRIMARY KEY)`
     await connection`CREATE TABLE workspace_files (id text PRIMARY KEY, workspace_id text REFERENCES workspace(id) ON DELETE CASCADE,
@@ -149,13 +186,9 @@ describe('chunked workspace file search on PostgreSQL', () => {
       '0313_puzzling_zodiak.sql',
       '0358_workspace_file_content_version_precision.sql',
       '0359_workspace_file_search_chunks.sql',
+      ginWriteMigration,
     ]) {
-      const source = readFileSync(
-        resolve(process.cwd(), '../../packages/db/migrations', migration),
-        'utf8'
-      ).replaceAll('"public".', `"${schema}".`)
-      for (const statement of source.split('--> statement-breakpoint'))
-        if (statement.trim()) await connection.unsafe(statement)
+      await applyMigration(migration)
     }
     database.current = drizzle(connection)
     database.search = drizzle(searchConnection, {
@@ -189,6 +222,88 @@ describe('chunked workspace file search on PostgreSQL', () => {
     }
   })
 
+  it('preserves search through disabling, draining, and replaying GIN pending-list maintenance', async () => {
+    await connection`ALTER INDEX workspace_file_search_chunk_content_idx SET (fastupdate = on)`
+    try {
+      await index('heading\nold needle αβγ\ntail')
+      const before = await ginState()
+      expect(before.pending_pages).toBeGreaterThan(0)
+
+      /** Simulate an interrupted rollout after the storage option commits but before the drain. */
+      await connection`ALTER INDEX workspace_file_search_chunk_content_idx SET (fastupdate = off)`
+      await index('heading\nnew needle αβγ\ntail', await addFile('file-2'))
+      expect((await ginState()).pending_pages).toBe(before.pending_pages)
+      const expected = [
+        { fileId: 'file-1', lineNumber: 2 },
+        { fileId: 'file-2', lineNumber: 2 },
+      ]
+      expect((await search('^(old|new) needle αβγ$', 'regex')).results).toMatchObject(expected)
+
+      for (let attempt = 0; attempt < 2; attempt++) {
+        await applyMigration(ginWriteMigration)
+        expect(await ginState()).toMatchObject({
+          oid: before.oid,
+          indisvalid: true,
+          direct_writes: true,
+          pending_pages: 0,
+        })
+        expect((await search('needle αβγ')).results).toMatchObject(expected)
+      }
+      await index('heading\nnew needle αβγ\ntail', await addFile('file-3'))
+      expect((await ginState()).pending_pages).toBe(0)
+      expect((await search('^(old|new) needle αβγ$', 'regex')).results).toHaveLength(3)
+    } finally {
+      await applyMigration(ginWriteMigration)
+    }
+  })
+
+  it('cancels a slow chunk statement without losing the connection or publishing partial content', async () => {
+    const build = (await beginFileSearchBuild(revision))!
+    const plan = planFileSearchIndex({ text: 'needle', partial: false }, signal)
+    const chunks = [...iterateFileSearchChunks(plan, signal)]
+    const writer = postgres(databaseUrl, {
+      max: 1,
+      prepare: false,
+      connection: { search_path: `${schema},public` },
+    })
+    const original = database.current
+    await connection`CREATE FUNCTION slow_chunk_insert() RETURNS trigger LANGUAGE plpgsql AS $$
+      BEGIN PERFORM pg_sleep(11); RETURN NEW; END $$`
+    await connection`CREATE TRIGGER slow_chunk_insert BEFORE INSERT ON workspace_file_search_chunk
+      FOR EACH STATEMENT EXECUTE FUNCTION slow_chunk_insert()`
+    try {
+      database.current = drizzle(writer)
+      const [before] = await writer`SELECT pg_backend_pid() AS pid`
+      await expect(appendFileSearchChunks(build, chunks, signal)).rejects.toMatchObject({
+        cause: { code: '57014' },
+      })
+      expect((await writer`SELECT pg_backend_pid() AS pid`)[0].pid).toBe(before.pid)
+      expect(
+        (await connection`SELECT count(*)::int AS count FROM workspace_file_search_chunk`)[0].count
+      ).toBe(0)
+      expect((await search('needle')).results).toEqual([])
+    } finally {
+      database.current = original
+      await writer.end()
+      await connection`DROP TRIGGER slow_chunk_insert ON workspace_file_search_chunk`
+      await connection`DROP FUNCTION slow_chunk_insert()`
+    }
+    expect(await appendFileSearchChunks(build, chunks, signal)).toBe(true)
+    expect(
+      await publishFileSearchBuild(
+        build,
+        {
+          status: 'ready',
+          chunkCount: chunks.length,
+          lineCount: plan.lineCount,
+          indexedBytes: plan.indexedBytes,
+        },
+        signal
+      )
+    ).toBe(true)
+    expect((await search('needle')).results).toMatchObject([{ fileId: 'file-1', lineNumber: 1 }])
+  })
+
   it('packs a million short lines without a million rows and bounds every stored value', async () => {
     await index('abc\n'.repeat(1_000_000))
     const [row] =

apps/sim/lib/workspace-files/search/constants.ts

@@ -46,7 +46,15 @@ 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
+/** Direct GIN writes perform index work in each insert, so transactions use smaller byte batches. */
+export const FILE_SEARCH_INSERT_BATCH_BYTES = 128 * 1024
+
+/** Index writes allow statement cancellation before the outer transaction terminates its session. */
+export const FILE_SEARCH_INDEX_TRANSACTION_LIMITS = {
+  statementTimeout: 10 * 1000,
+  lockTimeout: 5 * 1000,
+  transactionTimeout: 30 * 1000,
+} as const
 
 export const FILE_SEARCH_INDEX_GLOBAL_CONCURRENCY = 10
 export const FILE_SEARCH_INDEX_WORKSPACE_OUTSTANDING = 2

apps/sim/lib/workspace-files/search/index-state.ts

@@ -15,6 +15,7 @@ import {
   FILE_SEARCH_CLEANUP_BUDGET_MS,
   FILE_SEARCH_CLEANUP_MAX_BATCHES,
   FILE_SEARCH_CLEANUP_MIN_BATCH_MS,
+  FILE_SEARCH_INDEX_TRANSACTION_LIMITS,
   FILE_SEARCH_INSERT_BATCH_BYTES,
   FILE_SEARCH_INSERT_BATCH_ROWS,
 } from '@/lib/workspace-files/search/constants'
@@ -90,7 +91,7 @@ export async function beginFileSearchBuild(
   dispatchToken?: string
 ): Promise<FileSearchBuild | null> {
   return db.transaction(async (tx) => {
-    await configureFileSearchTransaction(tx)
+    await configureFileSearchTransaction(tx, FILE_SEARCH_INDEX_TRANSACTION_LIMITS)
     if (!(await lockCurrentFile(tx, revision))) return null
     const [observed] = await tx
       .select({
@@ -154,7 +155,7 @@ export async function appendFileSearchChunks(
     throw new Error('File search insert batch exceeds its budget')
   }
   return db.transaction(async (tx) => {
-    await configureFileSearchTransaction(tx)
+    await configureFileSearchTransaction(tx, FILE_SEARCH_INDEX_TRANSACTION_LIMITS)
     if (!(await lockBuild(tx, build))) return false
     signal.throwIfAborted()
     await tx
@@ -177,7 +178,7 @@ export async function publishFileSearchBuild(
   signal: AbortSignal
 ): Promise<boolean> {
   return db.transaction(async (tx) => {
-    await configureFileSearchTransaction(tx)
+    await configureFileSearchTransaction(tx, FILE_SEARCH_INDEX_TRANSACTION_LIMITS)
     signal.throwIfAborted()
     if (!(await lockCurrentFile(tx, build)) || !(await lockBuild(tx, build))) return false
     if (publication.status === 'ready') {
@@ -216,7 +217,7 @@ export async function failFileSearchRevision(
   dispatchToken?: string
 ): Promise<void> {
   await db.transaction(async (tx) => {
-    await configureFileSearchTransaction(tx)
+    await configureFileSearchTransaction(tx, FILE_SEARCH_INDEX_TRANSACTION_LIMITS)
     if (!(await lockCurrentFile(tx, revision))) return
     const [state] = await tx
       .select()

packages/db/migrations/0364_workspace_file_search_direct_gin.sql

@@ -0,0 +1,12 @@
+-- migration-safe: Metadata-only storage option; existing and new workers retain the same index and search semantics.
+ALTER INDEX "public"."workspace_file_search_chunk_content_idx" SET (fastupdate = off);
+--> statement-breakpoint
+-- Release the DDL lock before maintenance. Both operations are safe to replay after a partial migration.
+COMMIT;
+--> statement-breakpoint
+-- Disabling fastupdate does not flush existing pending entries. Drain them once without rebuilding the index.
+SET statement_timeout = '5min';
+--> statement-breakpoint
+SELECT pg_catalog.gin_clean_pending_list('"public"."workspace_file_search_chunk_content_idx"'::regclass);
+--> statement-breakpoint
+SET statement_timeout = 0;