fix(search): never grow a Search retirement page back to a size that timed out, and pause after a timeout

A timed-out page halved the row limit, but one fast page doubled it straight back, so the run
alternated between the size that timed out and half of it, rolling back a full statement-timeout
page each time. The limit now grows only up to half of the smallest size that timed out, and a
timed-out page is followed by the same pause as any other page.

Author: waleedlatif1

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

@@ -332,12 +332,15 @@ describe('retiring dormant Search embeddings', () => {
         CASE WHEN i % 3 = 0 THEN 'ordinary-doc' ELSE 'search-doc' END
       FROM generate_series(1003, 5002) i`
     await sql`CREATE TABLE committed_statement (rows integer NOT NULL)`
+    /** Sequences are not transactional, so this counts the timed-out statements that rolled back. */
+    await sql`CREATE SEQUENCE timed_out_statement`
     /** Stands in for write cost: a statement touching more than `bound` rows times out and rolls back. */
     await sql.unsafe(`CREATE FUNCTION bound_statement_rows() RETURNS trigger LANGUAGE plpgsql AS $$
       DECLARE touched integer;
       BEGIN
         SELECT count(*) INTO touched FROM changed_rows;
         IF touched > ${bound} THEN
+          PERFORM nextval('timed_out_statement');
           RAISE EXCEPTION 'canceling statement due to statement timeout' USING ERRCODE = 'query_canceled';
         END IF;
         INSERT INTO committed_statement VALUES (touched);
@@ -353,6 +356,13 @@ describe('retiring dormant Search embeddings', () => {
         sum(rows)::int AS total FROM committed_statement`
       expect(largest).toBeLessThanOrEqual(bound)
       expect(largest).toBeGreaterThan(0)
+      /**
+       * The limit halves 2,000 → 1,000 → 500 → 250 on the first document page and never grows back
+       * to a size that timed out, in either phase: exactly three rolled-back statements. A page
+       * that read past its row limit in either phase would time out again.
+       */
+      const [{ timeouts }] = await sql`SELECT last_value::int AS timeouts FROM timed_out_statement`
+      expect(timeouts).toBe(3)
       /**
        * Documents: i % 3 <> 0 gives 2,667 Search rows, of which i % 7 = 0 leaves 381 retired, so
        * 2,286 are updated, plus `search-doc`. Chunks: 501 Search rows from the fixture plus the
@@ -387,6 +397,7 @@ describe('retiring dormant Search embeddings', () => {
       await sql`DROP TRIGGER IF EXISTS bound_embedding_delete ON embedding`
       await sql`DROP FUNCTION bound_statement_rows()`
       await sql`DROP TABLE committed_statement`
+      await sql`DROP SEQUENCE timed_out_statement`
     }
   }, 60_000)
 
@@ -423,9 +434,9 @@ describe('retiring dormant Search embeddings', () => {
       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
+       * About 120 pages retire the 3,001 documents 25 at a time once the limit has halved down to
+       * its floor. A window of four IDs per row reads about 100 IDs and 25 update lookups per page,
+       * roughly 5 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)

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

@@ -23,9 +23,15 @@ const SCAN_ROWS_PER_MUTATION = 4
 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 doubles the row limit, widening its scan window with it. */
+/**
+ * A page faster than this doubles the row limit, widening its scan window with it, but never back
+ * to a size that timed out.
+ */
 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. */
+/**
+ * Each page, committed or timed out, 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
 const LOCK_RETRY_BUDGET_MS = 60_000
 
@@ -135,6 +141,8 @@ export const retireSearchEmbeddingsMigration: ScriptMigration = {
     let batches = 0
     let mutated = 0
     let rowLimit: number = ROW_LIMIT.initial
+    /** The largest limit the run may still try: half of the smallest limit that timed out. */
+    let ceiling: number = ROW_LIMIT.max
     for (;;) {
       /** Timed around the whole call, so the synchronous-replication wait at commit counts. */
       const pageStartedAt = performance.now()
@@ -145,8 +153,10 @@ export const retireSearchEmbeddingsMigration: ScriptMigration = {
         if (!(error instanceof PageMutationTimeout)) throw error
         /** The timed-out page rolled back with its cursor, so it is retried with fewer rows. */
         if (rowLimit <= ROW_LIMIT.min) throw error.timeout
-        rowLimit = halve(rowLimit)
+        ceiling = halve(rowLimit)
+        rowLimit = ceiling
         logger.warn('Search retirement page timed out; retrying with fewer rows', { rowLimit })
+        await sleep(Math.min(performance.now() - pageStartedAt, MAX_PAGE_PAUSE_MS))
         continue
       }
       if (page.done) break
@@ -167,7 +177,7 @@ export const retireSearchEmbeddingsMigration: ScriptMigration = {
           rowLimit,
         })
       } else if (pageMs < FAST_PAGE_MS) {
-        rowLimit = Math.min(ROW_LIMIT.max, rowLimit * 2)
+        rowLimit = Math.min(ceiling, rowLimit * 2)
       }
       if (batches % 10 === 0) {
         logger.info('Search embedding retirement progress', {

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

@@ -48,8 +48,9 @@ also widens the scan window, so sparse stretches are not crawled in small window
 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
 timeout. If a page's mutating statement exceeds the statement timeout, the page rolls back with its
-cursor and is retried with half the row limit; one that still times out at 25 rows fails the
-migration. Any other statement timeout fails the migration at once, because a smaller page cannot
+cursor and is retried with half the row limit after the usual pause. From then on, fast pages grow
+the limit only up to that halved size, so a size that timed out is never tried again. A page that
+still times out at 25 rows fails the migration. Any other statement timeout fails the migration at once, because a smaller page cannot
 speed it up. The completion rechecks, which walk every captured KB once, run with a 30-minute
 timeout. Brief lock timeouts retry the rolled-back page with bounded backoff for up to one minute.
 Other errors, or exhausted lock retries, fail the migration without a completion receipt.