fix(knowledge): serialize the Tin backfill with search-index membership changes

- Each backfill page takes its bases' membership locks, then reads under a fresh snapshot
- Remove a base's projection through the embedding knowledge base index instead of scanning the projection

Author: waleedlatif1

packages/db/script-migrations/0019_tin_keyword_projection.postgres.test.ts

@@ -1,6 +1,9 @@
 import { readFileSync } from 'node:fs'
 import path from 'node:path'
-import { installProjection } from '@sim/db/script-migrations/0019_tin_keyword_projection'
+import {
+  backfillProjection,
+  installProjection,
+} from '@sim/db/script-migrations/0019_tin_keyword_projection'
 import { generateId } from '@sim/utils/id'
 import postgres, { type Sql } from 'postgres'
 import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest'
@@ -58,6 +61,7 @@ describe.runIf(Boolean(databaseUrl))('Tin keyword projection in PostgreSQL', ()
       document_id text NOT NULL, enabled boolean NOT NULL DEFAULT true, content text,
       content_tsv tsvector NOT NULL
     )`
+    await sql`CREATE INDEX ON embedding (knowledge_base_id)`
     /** The shipped table, with its foreign key pointed at this schema's `embedding`. */
     const [createTable] = readFileSync(
       path.join(MIGRATIONS, '0366_tin_keyword_projection.sql'),
@@ -161,4 +165,51 @@ describe.runIf(Boolean(databaseUrl))('Tin keyword projection in PostgreSQL', ()
     await update
     expect(await projected()).toMatchObject([{ id: 'toggled', enabled: false }])
   })
+  it('removes a base through the chunk index instead of scanning the projection', async () => {
+    await sql`INSERT INTO embedding (id, knowledge_base_id, document_id, content_tsv)
+      SELECT 'index-' || n, 'index', 'doc', to_tsvector('english', 'Chunk ' || n)
+      FROM generate_series(1, 20) n`
+    /** Pending per-backend counts include earlier statements, so compare before and after. */
+    const scans = await sql.begin(async (tx) => {
+      const read = async () => {
+        const [row] = await tx<Array<{ scans: number }>>`
+          SELECT seq_scan::int AS scans FROM pg_stat_xact_user_tables
+          WHERE relid = 'embedding_keyword_tin'::regclass`
+        return row.scans
+      }
+      await tx`SET LOCAL enable_seqscan = off`
+      const before = await read()
+      await tx`UPDATE knowledge_base SET is_search_index = false WHERE id = 'index'`
+      return (await read()) - before
+    })
+    expect(scans).toBe(0)
+    expect(await projected()).toEqual([])
+  })
+
+  it('does not backfill a chunk whose base stops being a search index mid-page', async () => {
+    await sql`INSERT INTO embedding (id, knowledge_base_id, document_id, content_tsv)
+      VALUES ('unprojected', 'index', 'doc', to_tsvector('english', 'Predates the trigger'))`
+    await sql`DELETE FROM embedding_keyword_tin`
+    let demoted!: () => void
+    const demotedSignal = new Promise<void>((resolve) => {
+      demoted = resolve
+    })
+    let release!: () => void
+    const released = new Promise<void>((resolve) => {
+      release = resolve
+    })
+    const demotion = promoter.begin(async (tx) => {
+      await tx`UPDATE knowledge_base SET is_search_index = false WHERE id = 'index'`
+      demoted()
+      await released
+    })
+    await demotedSignal
+    const [{ pid }] = await sql`SELECT pg_backend_pid() AS pid`
+    const backfill = backfillProjection(sql)
+    await waitUntilBlocked(admin, pid)
+    release()
+    await demotion
+    expect(await backfill).toBe(0)
+    expect(await projected()).toEqual([])
+  })
 })

packages/db/script-migrations/0019_tin_keyword_projection.ts

@@ -56,7 +56,9 @@ async function tinAvailable(sql: Sql): Promise<boolean> {
  * each other, and a marker change takes it exclusively, so it waits for writers already in flight
  * and later writers wait for it. Each then reads under a fresh snapshot, so neither misses the
  * other's rows. A row lock cannot do this: a key-share lock on the marker row is compatible with
- * the uncommitted marker update, so a writer would read the old marker without waiting.
+ * the uncommitted marker update, so a writer would read the old marker without waiting. A whole
+ * base is reached through `embedding`'s knowledge base index, since a projected row always carries
+ * its chunk's base and the projection keeps no index but Tin's.
  */
 export async function installProjection(sql: Sql): Promise<void> {
   await sql.begin(async (tx) => {
@@ -102,7 +104,8 @@ export async function installProjection(sql: Sql): Promise<void> {
       BEGIN
         PERFORM pg_advisory_xact_lock(knowledge_tin_membership_key(NEW.id));
         IF NOT NEW.is_search_index THEN
-          DELETE FROM embedding_keyword_tin WHERE knowledge_base_id = NEW.id;
+          DELETE FROM embedding_keyword_tin t USING embedding e
+          WHERE e.knowledge_base_id = NEW.id AND t.id = e.id;
           RETURN NEW;
         END IF;
         INSERT INTO embedding_keyword_tin (id, knowledge_base_id, document_id, enabled, content)
@@ -122,40 +125,48 @@ export async function installProjection(sql: Sql): Promise<void> {
   })
 }
 
-/** Fills rows the trigger has not written yet, in independently committed keyset pages. */
-async function backfillProjection(sql: Sql): Promise<number> {
+/**
+ * Fills rows the trigger has not written yet, in independently committed keyset pages. Each page
+ * takes the membership locks of its bases before reading them, as the embedding trigger does, and
+ * then reads under a fresh snapshot: a base whose search-index marker is changing is read only
+ * after that change commits, so the backfill never writes back a chunk the change removed.
+ */
+export async function backfillProjection(sql: Sql): Promise<number> {
   let afterId = ''
   let inserted = 0
   let scanned = 0
   const startedAt = Date.now()
   for (;;) {
-    const [page] = await sql.begin(async (tx) => {
+    const page = await sql.begin(async (tx) => {
       await tx.unsafe("SET LOCAL lock_timeout = '5s'")
       await tx.unsafe("SET LOCAL statement_timeout = '60s'")
-      return tx.unsafe<Array<{ after_id: string | null; scanned: number; inserted: number }>>(
-        `
-        WITH source_page AS MATERIALIZED (
-          SELECT id FROM embedding WHERE id > $1 ORDER BY id LIMIT ${BATCH_SIZE}
-        ), batch AS MATERIALIZED (
+      const rows = await tx<Array<{ id: string; knowledge_base_id: string }>>`
+        SELECT id, knowledge_base_id FROM embedding WHERE id > ${afterId} ORDER BY id
+        LIMIT ${BATCH_SIZE}`
+      if (rows.length === 0) return null
+      const ids = rows.map((row) => row.id)
+      const bases = [...new Set(rows.map((row) => row.knowledge_base_id))]
+      await tx`SELECT pg_advisory_xact_lock_shared(knowledge_tin_membership_key(base))
+        FROM unnest(${bases}::text[]) AS base ORDER BY base`
+      const [{ written }] = await tx<Array<{ written: number }>>`
+        WITH batch AS MATERIALIZED (
           SELECT e.id, e.knowledge_base_id, e.document_id, e.enabled, e.content_tsv
-          FROM source_page p
-          INNER JOIN embedding e ON e.id = p.id
+          FROM embedding e
           INNER JOIN knowledge_base k ON k.id = e.knowledge_base_id AND k.is_search_index
-          WHERE NOT EXISTS (SELECT 1 FROM embedding_keyword_tin t WHERE t.id = p.id)
+          WHERE e.id = ANY(${ids}::text[])
+            AND NOT EXISTS (SELECT 1 FROM embedding_keyword_tin t WHERE t.id = e.id)
           ORDER BY e.id FOR KEY SHARE OF e
         ), written AS (
           INSERT INTO embedding_keyword_tin (id, knowledge_base_id, document_id, enabled, content)
           SELECT id, knowledge_base_id, document_id, enabled,
             knowledge_tin_base_token(knowledge_base_id) || ' ' || knowledge_tin_stream(content_tsv)
           FROM batch
           ON CONFLICT (id) DO NOTHING RETURNING id
-        ) SELECT max(id) AS after_id, count(*)::int AS scanned,
-          (SELECT count(*)::int FROM written) AS inserted FROM source_page`,
-        [afterId]
-      )
+        ) SELECT count(*)::int AS written FROM written`
+      return { afterId: ids[ids.length - 1], scanned: ids.length, inserted: written }
     })
-    if (!page.after_id) break
-    afterId = page.after_id
+    if (!page) break
+    afterId = page.afterId
     inserted += page.inserted
     scanned += page.scanned
     if (scanned % (BATCH_SIZE * 100) === 0) {