fix(search): lock each purge batch's rows before deleting them by tuple id

A concurrent update could move every chosen row to a new tuple id before the delete, so the batch deleted nothing and the purge reported the knowledge base finished with rows left. The batch now locks its rows in one statement and deletes them in the next, whose snapshot sees the current versions

Author: waleedlatif1

packages/db/scripts/purge-search-projection.integration.ts

@@ -8,6 +8,7 @@ import {
   SearchProjectionPurgeError,
 } from '@sim/db/scripts/purge-search-projection'
 import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure'
+import { sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
 import postgres, { type Sql } from 'postgres'
 import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest'
@@ -306,6 +307,31 @@ describe('Search candidate retirement in PostgreSQL', () => {
     expect(await candidateCount('search')).toBe(3)
   })
 
+  it('deletes rows a concurrent update moved after the batch chose them', async () => {
+    await seedBase({ id: 'search', width: 1536, model: 'text-embedding-3-small', chunks: 3 }, true)
+    const writer = postgres(fixtureUrl.toString(), { max: 1, onnotice: () => undefined })
+    try {
+      let purge: ReturnType<typeof purgeSearchProjection> | undefined
+      await writer.begin(async (tx) => {
+        // Every candidate row gets a new tuple version, committed only once the purge waits on it.
+        await tx`UPDATE embedding_search SET enabled = enabled WHERE knowledge_base_id = 'search'`
+        purge = purgeSearchProjection(sql, { batchSize: 10, pauseMs: 0 })
+        for (;;) {
+          const [{ waiting }] = await tx<{ waiting: boolean }[]>`
+            SELECT EXISTS (SELECT 1 FROM pg_stat_activity
+              WHERE datname = current_database() AND wait_event_type = 'Lock'
+                AND pid <> pg_backend_pid()) AS waiting`
+          if (waiting) break
+          await sleep(20)
+        }
+      })
+      expect(await purge).toEqual({ deleted: 3, batches: 1, remaining: false })
+    } finally {
+      await writer.end()
+    }
+    expect(await candidateCount('search')).toBe(0)
+  })
+
   it('stops when its connection is replaced, since the new session lacks the purge lock', async () => {
     await seedBase({ id: 'search', width: 1536, model: 'text-embedding-3-small', chunks: 5 }, true)
     let closed: () => void = () => undefined

packages/db/scripts/purge-search-projection.ts

@@ -202,8 +202,10 @@ interface SearchProjectionPurgeResult {
 
 /**
  * Deletes one batch of a Search knowledge base's candidate rows. The rows are chosen through
- * `embedding_search_kb_idx` and deleted by tuple id, re-checking the knowledge base so a row
- * moved or reassigned since the scan is left alone. Deletes nothing when the knowledge base is gone.
+ * `embedding_search_kb_idx` and locked by one statement, then deleted by tuple id in the next.
+ * The lock returns a concurrently updated row at its current tuple id and keeps it there, and the
+ * delete's later snapshot sees that version, so an empty batch means no row is left. The knowledge
+ * base re-check leaves a reassigned row alone. Deletes nothing when the knowledge base is gone.
  * Refuses when this backend no longer holds the purge lock: the client reconnects silently after a
  * dropped connection, and the new backend does not inherit the old one's session lock.
  */
@@ -224,13 +226,16 @@ async function purgeBatch(sql: Sql, knowledgeBaseId: string, batchSize: number):
         'A knowledge base is no longer Search-marked; stopped before deleting its candidate rows'
       )
     }
+    const [chosen] = await tx<{ ctids: string | null }[]>`
+      SELECT array_agg(ctid)::text AS ctids FROM (
+        SELECT ctid FROM embedding_search WHERE knowledge_base_id = ${knowledgeBaseId}
+        LIMIT ${batchSize} FOR UPDATE
+      ) batch`
+    if (!chosen.ctids) return 0
     const [row] = await tx<{ deleted: number }[]>`
       WITH deleted AS (
         DELETE FROM embedding_search
-        WHERE ctid = ANY (ARRAY(
-          SELECT ctid FROM embedding_search WHERE knowledge_base_id = ${knowledgeBaseId}
-          LIMIT ${batchSize}
-        )) AND knowledge_base_id = ${knowledgeBaseId}
+        WHERE ctid = ANY (${chosen.ctids}::tid[]) AND knowledge_base_id = ${knowledgeBaseId}
         RETURNING 1
       )
       SELECT count(*)::int AS deleted FROM deleted`