Skip to content

Commit 401cdee

Browse files
committed
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.
1 parent 60a1f6f commit 401cdee

3 files changed

Lines changed: 33 additions & 2 deletions

File tree

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -586,7 +586,7 @@ describe('retiring dormant Search embeddings', () => {
586586
await holding
587587
await blocker.end()
588588
}
589-
})
589+
}, 60_000)
590590

591591
it('maintains an already-retired index and resumes failed vacuum bookkeeping without rebuilding it again', async () => {
592592
await pass()
@@ -735,4 +735,21 @@ describe('retiring dormant Search embeddings', () => {
735735
await sql`DROP TABLE delete_page_started`
736736
}
737737
}, 60_000)
738+
739+
it('refuses to start while another retirement run holds the lock, leaving progress untouched', async () => {
740+
const other = postgres(readTestDatabaseUrl(), { max: 1, onnotice: () => undefined })
741+
try {
742+
await other`SELECT pg_advisory_lock(hashtextextended('search-embedding-retirement', 0))`
743+
await expect(retireSearchEmbeddings(sql)).rejects.toThrow('already running')
744+
expect(
745+
(await sql`SELECT to_regclass('search_embedding_cleanup_progress') AS relation`)[0].relation
746+
).toBeNull()
747+
} finally {
748+
await other.end()
749+
}
750+
await retireSearchEmbeddings(sql)
751+
expect(
752+
(await sql`SELECT count(*)::int AS n FROM embedding WHERE knowledge_base_id = 'search'`)[0].n
753+
).toBe(0)
754+
})
738755
})

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ const FAST_PAGE_MS = SLOW_PAGE_MS / 4
3333
/** The longest pause after one page, however slow the page was. */
3434
const MAX_PAGE_PAUSE_MS = 60_000
3535
const LOCK_RETRY_BUDGET_MS = 60_000
36+
const RETIREMENT_LOCK = 'search-embedding-retirement'
3637

3738
/**
3839
* 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(
118119
`Search retirement pacing needs a pause ratio of at least 0 and ${ROW_LIMIT.min}-${ROW_LIMIT.max} max rows`
119120
)
120121
}
122+
/** Session-level, like maintenance's lock, so overlapping operator runs never double the load. */
123+
const [{ locked }] =
124+
await sql`SELECT pg_try_advisory_lock(hashtextextended(${RETIREMENT_LOCK}, 0)) AS locked`
125+
if (!locked) throw new Error('Search retirement is already running')
126+
try {
127+
await retireTargets(sql, pacing)
128+
} finally {
129+
await sql`SELECT pg_advisory_unlock(hashtextextended(${RETIREMENT_LOCK}, 0))`
130+
}
131+
}
132+
133+
async function retireTargets(sql: Sql, pacing: RetirementPacing): Promise<void> {
121134
const pause = (pageMs: number) => sleep(Math.min(pageMs * pacing.pauseRatio, MAX_PAGE_PAUSE_MS))
122135
const hasTargets = await sql.begin('isolation level repeatable read', async (tx) => {
123136
await tx`SET LOCAL statement_timeout = '120s'`

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,8 @@ MIGRATION_DATABASE_URL=<direct DSN> bun run packages/db/script-migrations/0027_r
4343
Run it as the migration role: maintenance needs `pg_maintain`, which the application roles lack. Run
4444
it outside peak traffic, and run `--maintenance` in the quietest window you have: concurrent HNSW
4545
rebuilds are long and write a lot of WAL (GitLab, for example, schedules automatic reindexing for
46-
weekends). Keep one run at a time.
46+
weekends). One run at a time: a second run refuses to start while another holds the retirement
47+
lock, and maintenance has its own lock.
4748

4849
**Pausing.** Ctrl-C is safe at any point. The in-flight page rolls back with its cursor, and an
4950
interrupted concurrent rebuild's leftover index is removed on the next run. Rerun the same command to

0 commit comments

Comments
 (0)