From 401cdeec98b628f61809edb40c3f8886719b0ea1 Mon Sep 17 00:00:00 2001
From: Waleed Latif <walif6@gmail.com>
Date: Thu, 1 Oct 2026 10:33:50 -0700
Subject: [PATCH] fix(search): refuse overlapping Search retirement runs with a
 session lock

A plain operator run took no lock, so two runs could double the load on the primary. Retirement
now takes a session try-lock, as maintenance already does, and refuses to start while another run
holds it.
---
 ...27_retire_search_embeddings.integration.ts | 19 ++++++++++++++++++-
 .../0027_retire_search_embeddings.ts          | 13 +++++++++++++
 .../search-embedding-retirement.md            |  3 ++-
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/packages/db/script-migrations/0027_retire_search_embeddings.integration.ts b/packages/db/script-migrations/0027_retire_search_embeddings.integration.ts
index 34738fd2c78..bc4d47d7a0b 100644
--- a/packages/db/script-migrations/0027_retire_search_embeddings.integration.ts
+++ b/packages/db/script-migrations/0027_retire_search_embeddings.integration.ts
@@ -586,7 +586,7 @@ describe('retiring dormant Search embeddings', () => {
       await holding
       await blocker.end()
     }
-  })
+  }, 60_000)
 
   it('maintains an already-retired index and resumes failed vacuum bookkeeping without rebuilding it again', async () => {
     await pass()
@@ -735,4 +735,21 @@ describe('retiring dormant Search embeddings', () => {
       await sql`DROP TABLE delete_page_started`
     }
   }, 60_000)
+
+  it('refuses to start while another retirement run holds the lock, leaving progress untouched', async () => {
+    const other = postgres(readTestDatabaseUrl(), { max: 1, onnotice: () => undefined })
+    try {
+      await other`SELECT pg_advisory_lock(hashtextextended('search-embedding-retirement', 0))`
+      await expect(retireSearchEmbeddings(sql)).rejects.toThrow('already running')
+      expect(
+        (await sql`SELECT to_regclass('search_embedding_cleanup_progress') AS relation`)[0].relation
+      ).toBeNull()
+    } finally {
+      await other.end()
+    }
+    await retireSearchEmbeddings(sql)
+    expect(
+      (await sql`SELECT count(*)::int AS n FROM embedding WHERE knowledge_base_id = 'search'`)[0].n
+    ).toBe(0)
+  })
 })
diff --git a/packages/db/script-migrations/0027_retire_search_embeddings.ts b/packages/db/script-migrations/0027_retire_search_embeddings.ts
index f58f5b14b01..9c3df7f5f56 100644
--- a/packages/db/script-migrations/0027_retire_search_embeddings.ts
+++ b/packages/db/script-migrations/0027_retire_search_embeddings.ts
@@ -33,6 +33,7 @@ const FAST_PAGE_MS = SLOW_PAGE_MS / 4
 /** The longest pause after one page, however slow the page was. */
 const MAX_PAGE_PAUSE_MS = 60_000
 const LOCK_RETRY_BUDGET_MS = 60_000
+const RETIREMENT_LOCK = 'search-embedding-retirement'
 
 /**
  * How hard one run pushes the primary. Each page, committed or timed out, is followed by a pause of
@@ -118,6 +119,18 @@ export async function retireSearchEmbeddings(
       `Search retirement pacing needs a pause ratio of at least 0 and ${ROW_LIMIT.min}-${ROW_LIMIT.max} max rows`
     )
   }
+  /** Session-level, like maintenance's lock, so overlapping operator runs never double the load. */
+  const [{ locked }] =
+    await sql`SELECT pg_try_advisory_lock(hashtextextended(${RETIREMENT_LOCK}, 0)) AS locked`
+  if (!locked) throw new Error('Search retirement is already running')
+  try {
+    await retireTargets(sql, pacing)
+  } finally {
+    await sql`SELECT pg_advisory_unlock(hashtextextended(${RETIREMENT_LOCK}, 0))`
+  }
+}
+
+async function retireTargets(sql: Sql, pacing: RetirementPacing): Promise<void> {
   const pause = (pageMs: number) => sleep(Math.min(pageMs * pacing.pauseRatio, MAX_PAGE_PAUSE_MS))
   const hasTargets = await sql.begin('isolation level repeatable read', async (tx) => {
     await tx`SET LOCAL statement_timeout = '120s'`
diff --git a/packages/db/script-migrations/search-embedding-retirement.md b/packages/db/script-migrations/search-embedding-retirement.md
index 86d48536e5b..26a8ce59bfd 100644
--- a/packages/db/script-migrations/search-embedding-retirement.md
+++ b/packages/db/script-migrations/search-embedding-retirement.md
@@ -43,7 +43,8 @@ MIGRATION_DATABASE_URL=<direct DSN> bun run packages/db/script-migrations/0027_r
 Run it as the migration role: maintenance needs `pg_maintain`, which the application roles lack. Run
 it outside peak traffic, and run `--maintenance` in the quietest window you have: concurrent HNSW
 rebuilds are long and write a lot of WAL (GitLab, for example, schedules automatic reindexing for
-weekends). Keep one run at a time.
+weekends). One run at a time: a second run refuses to start while another holds the retirement
+lock, and maintenance has its own lock.
 
 **Pausing.** Ctrl-C is safe at any point. The in-flight page rolls back with its cursor, and an
 interrupted concurrent rebuild's leftover index is removed on the next run. Rerun the same command to