fix(search): scale the retirement scan window with the row limit and time out resumed rechecks like completion

A capped page re-reads its scan from its last mutated row, so a fixed 25,000-ID window re-read
most of the same IDs on every page once the row limit shrank. Each page now reads at most four IDs
per row of its limit, capped at 25,000, and any fast page doubles the limit so sparse stretches
widen the window again.

A retry that finds retirement already complete revalidates every captured KB, as completion does,
so it now runs under the same 30-minute timeout instead of the two-minute page timeout.

Author: waleedlatif1

packages/db/script-migrations/0027_retire_search_embeddings.integration.ts

@@ -390,6 +390,84 @@ describe('retiring dormant Search embeddings', () => {
     }
   }, 60_000)
 
+  it('bounds the IDs each page reads by the row limit once the limit shrinks', async () => {
+    const docs = 3000
+    const bound = 25
+    await sql`INSERT INTO document (id, knowledge_base_id)
+      SELECT 'doc-' || lpad(i::text, 5, '0'), 'search' FROM generate_series(1, ${docs}) i`
+    /** Statements over `bound` rows time out, which pins the row limit at its 25-row floor. */
+    await sql.unsafe(`CREATE FUNCTION bound_document_update() RETURNS trigger LANGUAGE plpgsql AS $$
+      BEGIN
+        IF (SELECT count(*) FROM changed_rows) > ${bound} THEN
+          RAISE EXCEPTION 'canceling statement due to statement timeout' USING ERRCODE = 'query_canceled';
+        END IF;
+        RETURN NULL;
+      END $$`)
+    await sql`CREATE TRIGGER bound_document_update AFTER UPDATE ON document
+      REFERENCING NEW TABLE AS changed_rows FOR EACH STATEMENT EXECUTE FUNCTION bound_document_update()`
+    /**
+     * On a table this small the planner may answer any page with a sequential scan, which reads
+     * every row whatever the window; production pages use the primary key, so the test does too.
+     */
+    await sql`SET enable_seqscan = off`
+    /** Document rows read by any scan, counted across committed and rolled-back pages alike. */
+    async function documentReads() {
+      await sql`SELECT pg_stat_force_next_flush()`
+      await admin`SELECT pg_stat_clear_snapshot()`
+      const [row] = await admin`SELECT (seq_tup_read + coalesce(idx_tup_fetch, 0))::int AS n
+        FROM pg_stat_user_tables WHERE schemaname = ${schema} AND relname = 'document'`
+      return row.n
+    }
+    try {
+      const before = await documentReads()
+      expect(await pass()).toBe(true)
+      const reads = (await documentReads()) - before
+      /**
+       * About 120 pages retire the 3,001 documents 25 at a time, each after one rolled-back attempt
+       * at 50 rows. A window of four IDs per row reads about 300 IDs and 75 update lookups per page,
+       * roughly 15 reads per document, plus the chunk phase and completion rechecks. A fixed
+       * 25,000-ID window re-reads the rest of the table on every attempt, over 100 per document.
+       */
+      expect(reads).toBeLessThan(30 * docs)
+      expect(
+        (
+          await sql`SELECT count(*)::int AS n FROM document
+          WHERE knowledge_base_id = 'search' AND (NOT user_excluded OR enabled)`
+        )[0].n
+      ).toBe(0)
+    } finally {
+      await sql`RESET enable_seqscan`
+      await sql`DROP TRIGGER IF EXISTS bound_document_update ON document`
+      await sql`DROP FUNCTION bound_document_update()`
+    }
+  }, 120_000)
+
+  it('gives the completed-retirement recheck the completion timeout when it resumes', async () => {
+    expect(await pass()).toBe(true)
+    await sql`DELETE FROM script_migrations`
+    /** Stands in for a recheck that outlasts the two-minute page timeout on a large target set. */
+    await sql`CREATE FUNCTION require_recheck_timeout() RETURNS boolean LANGUAGE plpgsql AS $$
+      BEGIN
+        IF current_setting('statement_timeout') <> '30min' THEN
+          RAISE EXCEPTION 'canceling statement due to statement timeout' USING ERRCODE = 'query_canceled';
+        END IF;
+        RETURN true;
+      END $$`
+    await sql`ALTER TABLE search_embedding_cleanup_targets RENAME TO captured_targets`
+    await sql`CREATE VIEW search_embedding_cleanup_targets AS
+      SELECT knowledge_base_id FROM captured_targets WHERE require_recheck_timeout()`
+    try {
+      expect(await sql`SELECT phase FROM search_embedding_cleanup_progress`).toEqual([
+        { phase: 'done' },
+      ])
+      expect(await pass()).toBe(true)
+    } finally {
+      await sql`DROP VIEW IF EXISTS search_embedding_cleanup_targets`
+      await sql`ALTER TABLE IF EXISTS captured_targets RENAME TO search_embedding_cleanup_targets`
+      await sql`DROP FUNCTION require_recheck_timeout()`
+    }
+  })
+
   it('scans past a run of already-retired documents longer than the row limit in one page', async () => {
     /** 6,000 retired Search documents exceed the initial 2,000-row limit but fit one 25,000-ID scan. */
     await sql`INSERT INTO document (id, knowledge_base_id, user_excluded, enabled)

packages/db/script-migrations/0027_retire_search_embeddings.ts

@@ -7,8 +7,14 @@ import { sleep } from '@sim/utils/helpers'
 import postgres, { type Sql, type TransactionSql } from 'postgres'
 
 const logger = createLogger('RetireSearchEmbeddings')
-/** IDs one page reads in primary-key order; reading is cheap next to the mutation. */
+/** Most IDs one page reads in primary-key order; reading is cheap next to the mutation. */
 const SCAN_PAGE_SIZE = 25_000
+/**
+ * A page reads at most this many IDs per row it may mutate. A page that reaches the row limit is
+ * re-read from its last mutated row, so a window far wider than the limit would be read again and
+ * again as the limit shrinks, exactly when the database is already slow.
+ */
+const SCAN_ROWS_PER_MUTATION = 4
 /**
  * Rows one page may update or delete. Every retired document is a non-HOT update touching each of
  * its indexes, and every deleted chunk cascades into its projections, so the write cost of a page,
@@ -17,7 +23,7 @@ const SCAN_PAGE_SIZE = 25_000
 const ROW_LIMIT = { initial: 2_000, min: 25, max: 8_000 } as const
 /** A page slower than this halves the row limit. */
 const SLOW_PAGE_MS = 30_000
-/** A page faster than this that reached the row limit doubles it. */
+/** A page faster than this doubles the row limit, widening its scan window with it. */
 const FAST_PAGE_MS = SLOW_PAGE_MS / 4
 /** Each page is followed by a pause as long as the page, up to this, to leave the primary headroom. */
 const MAX_PAGE_PAUSE_MS = 5_000
@@ -33,8 +39,6 @@ interface PageResult {
   afterId: string
   /** Rows the page updated or deleted. */
   mutated: number
-  /** The page stopped at the row limit rather than the end of its scan. */
-  limited: boolean
   /** A phase change the page committed. */
   transition?: 'embeddings' | 'documents_rescan' | 'embeddings_rescan'
 }
@@ -162,7 +166,7 @@ export const retireSearchEmbeddingsMigration: ScriptMigration = {
           pageMs: Math.round(pageMs),
           rowLimit,
         })
-      } else if (page.limited && pageMs < FAST_PAGE_MS) {
+      } else if (pageMs < FAST_PAGE_MS) {
         rowLimit = Math.min(ROW_LIMIT.max, rowLimit * 2)
       }
       if (batches % 10 === 0) {
@@ -186,14 +190,16 @@ export const retireSearchEmbeddingsMigration: ScriptMigration = {
 }
 
 /**
- * Retires one page: the target rows among the next `SCAN_PAGE_SIZE` IDs, capped at `rowLimit`.
+ * Retires one page: the target rows among the next `rowLimit * SCAN_ROWS_PER_MUTATION` IDs (at
+ * most `SCAN_PAGE_SIZE`), capped at `rowLimit`.
  * A capped page advances the cursor only to its last mutated row, so the rest of the scan is
  * read again by the next page; an uncapped page advances past its whole scan.
  */
 async function retirePage(sql: Sql, rowLimit: number): Promise<PageResult> {
   return retryOnLockTimeout(
     () =>
       sql.begin(async (tx) => {
+        const scanLimit = Math.min(SCAN_PAGE_SIZE, rowLimit * SCAN_ROWS_PER_MUTATION)
         await tx`SET LOCAL statement_timeout = '120s'`
         await tx`SET LOCAL lock_timeout = '1s'`
         const [progress] = await tx<Progress[]>`
@@ -203,24 +209,25 @@ async function retirePage(sql: Sql, rowLimit: number): Promise<PageResult> {
           phase: progress.phase,
           afterId: progress.after_id,
           mutated: 0,
-          limited: false,
           ...page,
         })
         if (progress.phase === 'done') {
+          /** A retry after maintenance failed rechecks every captured KB, as completion did. */
+          await tx`SET LOCAL statement_timeout = '30min'`
           await validateTargetMarkers(tx)
           return result({ done: true })
         }
 
         if (progress.phase === 'documents') {
           /** Already-retired documents are skipped so they never spend the row limit. */
           const [page] = await pageMutation(tx<
-            { after_id: string; limited: boolean; mutated: number; invalid_target: boolean }[]
+            { after_id: string; mutated: number; invalid_target: boolean }[]
           >`
             WITH source_page AS MATERIALIZED (
               SELECT id, knowledge_base_id,
                 (NOT user_excluded OR enabled OR processing_queue_token IS NOT NULL
                  OR processing_queued_at IS NOT NULL OR processing_deferred_until IS NOT NULL) AS unretired
-              FROM document WHERE id > ${progress.after_id} ORDER BY id LIMIT ${SCAN_PAGE_SIZE}
+              FROM document WHERE id > ${progress.after_id} ORDER BY id LIMIT ${scanLimit}
             ), target_page AS MATERIALIZED (
               SELECT p.id, p.knowledge_base_id FROM source_page p
               JOIN search_embedding_cleanup_targets t ON t.knowledge_base_id = p.knowledge_base_id
@@ -246,7 +253,7 @@ async function retirePage(sql: Sql, rowLimit: number): Promise<PageResult> {
                      OR d.processing_queued_at IS NOT NULL OR d.processing_deferred_until IS NOT NULL)
               RETURNING d.id
             ) SELECT CASE WHEN e.limited THEN e.last_target ELSE max(p.id) END AS after_id,
-                e.limited, (SELECT count(*) FROM retired)::int AS mutated,
+                (SELECT count(*) FROM retired)::int AS mutated,
                 EXISTS (SELECT 1 FROM invalid_target) AS invalid_target
               FROM source_page p CROSS JOIN page_end e GROUP BY e.limited, e.last_target`)
           if (!page) {
@@ -256,21 +263,20 @@ async function retirePage(sql: Sql, rowLimit: number): Promise<PageResult> {
           if (page.invalid_target)
             throw new Error('Cleanup target is no longer a Search knowledge base')
           await tx`UPDATE search_embedding_cleanup_progress SET after_id = ${page.after_id} WHERE id = 1`
-          return result({ afterId: page.after_id, mutated: page.mutated, limited: page.limited })
+          return result({ afterId: page.after_id, mutated: page.mutated })
         }
 
         /** Keep page IDs in PostgreSQL; foreign keys cascade projection and provenance deletes. */
         const [page] = await pageMutation(tx<
           {
             after_id: string
-            limited: boolean
             mutated: number
             unretired: boolean
             invalid_target: boolean
           }[]
         >`
           WITH source_page AS MATERIALIZED (
-            SELECT id, knowledge_base_id, document_id FROM embedding WHERE id > ${progress.after_id} ORDER BY id LIMIT ${SCAN_PAGE_SIZE}
+            SELECT id, knowledge_base_id, document_id FROM embedding WHERE id > ${progress.after_id} ORDER BY id LIMIT ${scanLimit}
           ), target_page AS MATERIALIZED (
             SELECT p.* FROM source_page p
             JOIN search_embedding_cleanup_targets t ON t.knowledge_base_id = p.knowledge_base_id
@@ -295,7 +301,7 @@ async function retirePage(sql: Sql, rowLimit: number): Promise<PageResult> {
               AND NOT EXISTS (SELECT 1 FROM unretired) AND NOT EXISTS (SELECT 1 FROM invalid_target)
             RETURNING e.id
           ) SELECT CASE WHEN e.limited THEN e.last_target ELSE max(p.id) END AS after_id,
-              e.limited, (SELECT count(*) FROM deleted)::int AS mutated,
+              (SELECT count(*) FROM deleted)::int AS mutated,
               EXISTS (SELECT 1 FROM unretired) AS unretired,
               EXISTS (SELECT 1 FROM invalid_target) AS invalid_target
             FROM source_page p CROSS JOIN page_end e GROUP BY e.limited, e.last_target`)
@@ -329,7 +335,7 @@ async function retirePage(sql: Sql, rowLimit: number): Promise<PageResult> {
           return result({ done: true })
         }
         await tx`UPDATE search_embedding_cleanup_progress SET after_id = ${page.after_id} WHERE id = 1`
-        return result({ afterId: page.after_id, mutated: page.mutated, limited: page.limited })
+        return result({ afterId: page.after_id, mutated: page.mutated })
       }),
     {
       budgetMs: LOCK_RETRY_BUDGET_MS,

packages/db/script-migrations/search-embedding-retirement.md

@@ -30,18 +30,20 @@ resume the saved scope, phase, cursor and maintenance checkpoints.
 The existing maintenance implementation rebuilds HNSW indexes and vacuums affected tables before
 deployment continues.
 
-Each page reads at most 25,000 IDs and mutates at most a row limit of them. Pages execute
+Each page mutates at most a row limit of target rows and reads at most four IDs per row of that
+limit, never more than 25,000 IDs. Pages execute
 sequentially, and each is followed by a pause as long as the page took, up to five seconds, to
 leave the primary headroom. Retiring a document is a non-HOT update that writes every index on
 `document`, and deleting a chunk cascades into its projections, so a page's cost follows the target
 rows it mutates, not the IDs it reads. A page that reaches the row limit advances the cursor only to
-its last mutated row; the rest of its scan is read again by the next page. Documents that are
-already retired never count against the limit.
+its last mutated row; the rest of its scan is read again by the next page. Tying the scan window to
+the limit keeps that re-reading proportional to the work, even after the limit shrinks. Documents
+that are already retired never count against the limit.
 
 The row limit starts at 2,000 rows. A page is timed from the start of its transaction through its
 commit, including the synchronous-replication wait and any lock-timeout retries. A page slower than
-30 seconds halves the limit. A fast page, one under 7.5 seconds, that reached the limit doubles it,
-up to 8,000. The limit never drops below 25 rows. Phase changes do not adjust it.
+30 seconds halves the limit. A fast page, one under 7.5 seconds, doubles it up to 8,000, which
+also widens the scan window, so sparse stretches are not crawled in small windows. The limit never drops below 25 rows. Phase changes do not adjust it.
 
 Materialized SQL pages keep the IDs inside PostgreSQL; the migration process receives only a cursor
 and a validation result. Each page uses a two-minute statement timeout and a one-second lock
@@ -88,7 +90,7 @@ unrelated rows. Before completion, the cleanup checks for unretired documents an
 behind either cursor and restarts the affected phase if needed. A final bounded pass validates all
 captured KB markers, including empty KBs and KBs whose rows were already scanned, holding shared
 marker locks until the completion checkpoint commits. Resuming a completed cleanup before maintenance
-also revalidates the captured set. Keep target writers stopped and
+also revalidates the captured set, with the same 30-minute timeout as the completion recheck. Keep target writers stopped and
 do not change their Search markers during the pass.
 
 Inspect progress with: