From 66516f52d2230278a6e016f3cba3ae3ffd75c53d Mon Sep 17 00:00:00 2001 From: Garry Tan Date: Tue, 16 Jun 2026 22:08:42 -0700 Subject: [PATCH] feat(embed): wire DB-pacing into embed paths + single-flight + bounded keyset re-entry embedStaleForSource + CLI embedAllStale/embedAll lower worker count to the resolved cap and observe()/pace() their DB ops; embed job + embed-backfill handler resolve env>config>bundle; CLI --pace flags; --background carries overrides into the job payload; single-flight via shared per-source lock; budget-timer re-arm around paced sleeps; EmbedResult.pacing telemetry. --- src/commands/embed.ts | 326 ++++++++++++++++++-- src/commands/jobs.ts | 8 + src/core/embed-stale.ts | 47 ++- src/core/minions/handlers/embed-backfill.ts | 45 ++- 4 files changed, 386 insertions(+), 40 deletions(-) diff --git a/src/commands/embed.ts b/src/commands/embed.ts index 3afefab03..7c324c588 100644 --- a/src/commands/embed.ts +++ b/src/commands/embed.ts @@ -9,7 +9,16 @@ import { loadConfig } from '../core/config.ts'; import { slog, serr } from '../core/console-prefix.ts'; import { filterOutEmbedSkipped } from '../core/embed-skip.ts'; import { runSlidingPool } from '../core/worker-pool.ts'; -import { isAborted, anySignal } from '../core/abort-check.ts'; +import { isAborted, anySignal, AbortError } from '../core/abort-check.ts'; +import { type DbPacer, createDbPacer, createNoopPacer, observed } from '../core/db-pacer.ts'; +import { + resolvePaceMode, + loadPaceModeConfig, + readPaceEnv, + type PaceKeyOverrides, +} from '../core/pace-mode.ts'; +import { tryAcquireDbLock, type DbLockHandle } from '../core/db-lock.ts'; +import { embedBackfillLockId } from '../core/embed-backfill-lock.ts'; export interface EmbedOpts { /** Embed ALL pages (every chunk). */ @@ -71,6 +80,26 @@ export interface EmbedOpts { * with the internal wall-clock budget timer via `anySignal`. */ signal?: AbortSignal; + /** + * DB-contention pacing (paced-backfill). Raw inputs resolved in + * runEmbedCore via env > config > bundle (env beats config = incident + * escape hatch). `perCallMode` is from `--pace[=mode]`; `perCall` from + * `--pace-max-concurrency` etc. Absent ⇒ resolves from env/config (so a + * queued job paced by config alone still throttles). Mode `off` ⇒ no-op. + */ + pace?: { + perCallMode?: string; + perCall?: PaceKeyOverrides; + }; + /** + * E-2 (paced-backfill): single-flight the stale run by taking the SAME + * per-source lock the `embed-backfill` minion handler uses, so a hand-run CLI + * backfill and a queued job can't grind the same source at once (closing the + * NULL→non-NULL upsert race window that paced — longer — runs widen). Set + * ONLY by the CLI (`runEmbed`); the minion path already locks. All-source + * runs lock every source in sorted order. dryRun skips it. + */ + singleFlight?: boolean; } /** @@ -94,6 +123,24 @@ export interface EmbedResult { pages_processed: number; /** True if this run was a dry-run. */ dryRun: boolean; + /** + * E1 (paced-backfill): end-of-run pacing telemetry. Present ONLY when pacing + * was active (enabled bundle). The number the operator could not get from an + * external wrapper ("zero pauses" ≠ "queue safe"). + */ + pacing?: { + maxConcurrency: number; + /** In-band latency samples folded into the EWMA. */ + samples: number; + /** Final EWMA of observed DB-op latency (ms), or null if no samples. */ + ewmaMs: number | null; + /** Cumulative cooperative-sleep time (ms). */ + totalSleptMs: number; + /** Number of cooperative sleeps. */ + sleeps: number; + /** High-water mark of acquirers blocked on the permit (sync path). */ + maxWaiters: number; + }; } /** @@ -207,11 +254,110 @@ export async function runEmbedCore(engine: BrainEngine, opts: EmbedOpts): Promis return result; } if (opts.all || opts.stale) { - await embedAll(engine, !!opts.stale, !!opts.dryRun, result, opts.onProgress, opts.sourceId, { - batchSize: opts.batchSize, - priority: opts.priority, - catchUp: opts.catchUp, - }, opts.signal); + // E-2 (paced-backfill): CLI single-flight. Take the SAME per-source lock as + // the embed-backfill minion handler so a hand-run backfill and a queued job + // are mutually exclusive per source. All-source runs lock every source in + // sorted (deterministic) order to avoid acquire-order deadlock. Released in + // the finally below. Skipped for dryRun and when the caller didn't opt in + // (cycle / catch-up / sync-auto-embed callers never single-flight). + const sfLocks: DbLockHandle[] = []; + if (opts.singleFlight && opts.stale && !opts.dryRun) { + let lockSourceIds: string[]; + if (opts.sourceId) { + lockSourceIds = [opts.sourceId]; + } else { + try { + const rows = await engine.listAllSources(); + lockSourceIds = rows.map((r) => r.id).sort(); + } catch { + lockSourceIds = []; + } + } + for (const sid of lockSourceIds) { + let lock: DbLockHandle | null = null; + try { + lock = await tryAcquireDbLock(engine, embedBackfillLockId(sid), 60); + } catch { + // Fail-open: a lock-subsystem error must not crash a backfill. Drop + // single-flight for this run (release what we took) and proceed. + for (const h of sfLocks) { + try { await h.release(); } catch { /* best-effort */ } + } + sfLocks.length = 0; + break; + } + if (!lock) { + // Another backfill (CLI or job) holds this source. Release what we + // took and bail cleanly rather than racing the upsert path. + for (const h of sfLocks) { + try { await h.release(); } catch { /* best-effort */ } + } + serr(` [embed] another backfill is already running for source "${sid}"; skipping (single-flight).`); + return result; + } + sfLocks.push(lock); + } + } + + // Resolve DB-contention pacing (env > config > bundle; env is the + // incident escape hatch). dryRun skips it — no writes to pace. A + // disabled bundle yields a no-op pacer (zero overhead on the hot path). + let pacer: DbPacer = createNoopPacer(); + let paceMaxConcurrency: number | undefined; + if (!opts.dryRun) { + try { + const cfg = await loadPaceModeConfig(engine); + const { envMode, envOverrides } = readPaceEnv(); + const knobs = resolvePaceMode({ + mode: cfg.mode, + configOverrides: cfg.configOverrides, + envMode, + envOverrides, + perCallMode: opts.pace?.perCallMode, + perCall: opts.pace?.perCall, + }); + if (knobs.enabled) { + pacer = createDbPacer({ bundle: knobs }); + paceMaxConcurrency = knobs.maxConcurrency; + } + } catch { + // Fail-open: pacing must never break a backfill. + pacer = createNoopPacer(); + } + } + try { + await embedAll(engine, !!opts.stale, !!opts.dryRun, result, opts.onProgress, opts.sourceId, { + batchSize: opts.batchSize, + priority: opts.priority, + catchUp: opts.catchUp, + pacer, + paceMaxConcurrency, + }, opts.signal); + } finally { + // E1: surface pacing telemetry (human + structured) when pacing was on. + const snap = pacer.snapshot(); + if (snap.enabled) { + result.pacing = { + maxConcurrency: snap.maxConcurrency, + samples: snap.sampleCount, + ewmaMs: snap.ewmaMs, + totalSleptMs: snap.totalSleptMs, + sleeps: snap.sleepCount, + maxWaiters: snap.maxWaiters, + }; + serr( + ` [embed] pacing: cap=${snap.maxConcurrency} samples=${snap.sampleCount} ` + + `ewma=${snap.ewmaMs === null ? 'n/a' : Math.round(snap.ewmaMs) + 'ms'} ` + + `slept=${snap.totalSleptMs}ms/${snap.sleepCount}`, + ); + } + pacer.dispose(); + // E-2: release single-flight locks (reverse order). Best-effort; the + // lock TTL is the backstop if a release fails. + for (const h of sfLocks.reverse()) { + try { await h.release(); } catch { /* best-effort; TTL covers it */ } + } + } return result; } if (opts.slug) { @@ -221,6 +367,38 @@ export async function runEmbedCore(engine: BrainEngine, opts: EmbedOpts): Promis throw new Error('No embed target specified. Pass { slug }, { slugs }, { all }, or { stale }.'); } +/** + * Parse the `--pace` family from a CLI arg list. Returns ONLY the explicit + * overrides (CX5: never the full resolved bundle) so they can be serialized + * into a background-job payload and re-resolved (env > config > bundle) at + * execution. Returns undefined when no pace flag is present. + * + * Recognized: `--pace` (bare ⇒ balanced), `--pace=`, + * `--pace-max-concurrency=` / `--pace-max-concurrency `. + */ +export function parsePaceArgs( + args: string[], +): { perCallMode?: string; perCall?: PaceKeyOverrides } | undefined { + let perCallMode: string | undefined; + let perCall: PaceKeyOverrides | undefined; + for (let i = 0; i < args.length; i++) { + const a = args[i]; + if (a === '--pace') { + perCallMode = 'balanced'; + } else if (a.startsWith('--pace=')) { + perCallMode = a.slice('--pace='.length) || 'balanced'; + } else if (a.startsWith('--pace-max-concurrency=')) { + const n = parseInt(a.slice('--pace-max-concurrency='.length), 10); + if (Number.isFinite(n) && n >= 1) (perCall ??= {}).maxConcurrency = n; + } else if (a === '--pace-max-concurrency') { + const n = parseInt(args[i + 1] ?? '', 10); + if (Number.isFinite(n) && n >= 1) (perCall ??= {}).maxConcurrency = n; + } + } + if (perCallMode === undefined && perCall === undefined) return undefined; + return { ...(perCallMode !== undefined && { perCallMode }), ...(perCall && { perCall }) }; +} + export async function runEmbed(engine: BrainEngine, args: string[]): Promise { // v0.36+ T7: --background submits via Minion queue, returns job_id to // stdout, exits. Same semantics in TTY and cron (D9). @@ -239,6 +417,10 @@ export async function runEmbed(engine: BrainEngine, args: string[]): Promise= 0 ? cleanArgs.slice(slugsI + 1).filter(a => !a.startsWith('--')) : undefined, sourceId: srcI >= 0 ? cleanArgs[srcI + 1] : undefined, + // CX1+CX5: carry explicit pace overrides into the `embed` job payload + // (the job name CLI --background actually submits). The handler + // re-resolves env > config > bundle at execution. + ...(parsePaceArgs(cleanArgs) && { pace: parsePaceArgs(cleanArgs) }), }; }, source: 'cli', @@ -262,12 +444,14 @@ export async function runEmbed(engine: BrainEngine, args: string[]): Promise= 0 ? args[priorityIdx + 1] : undefined; const priority = priorityRaw === 'recent' ? 'recent' as const : undefined; const catchUp = args.includes('--catch-up'); + const pace = parsePaceArgs(args); let opts: EmbedOpts; if (slugsIdx >= 0) { opts = { slugs: args.slice(slugsIdx + 1).filter(a => !a.startsWith('--')), dryRun, sourceId, batchSize, priority, catchUp }; } else if (all || stale) { - opts = { all, stale, dryRun, sourceId, batchSize, priority, catchUp }; + // E-2: CLI-only single-flight for stale runs (the minion path locks itself). + opts = { all, stale, dryRun, sourceId, batchSize, priority, catchUp, ...(pace && { pace }), ...(stale && { singleFlight: true }) }; } else { const slug = args.find(a => !a.startsWith('--')); if (!slug) { @@ -416,6 +600,10 @@ async function embedAll( batchSize?: number; priority?: 'recent'; catchUp?: boolean; + /** DB-contention pacer (paced-backfill); no-op when pacing is off. */ + pacer?: DbPacer; + /** Resolved concurrency cap (E-1: the worker count, no separate permit). */ + paceMaxConcurrency?: number; }, signal?: AbortSignal, ) { @@ -443,6 +631,10 @@ async function embedAll( return await embedAllStale(engine, sourceId, dryRun, result, onProgress, staleOpts, signature, signal); } + // --all path: pacer (no-op when off). E-1: lower the worker count to the + // resolved cap instead of adding a separate permit. + const pacer = staleOpts?.pacer ?? createNoopPacer(); + // v0.31.12: when sourceId is set, scope listPages to that source. // v0.41 (D8 + Codex r2 #11): apply embed-skip filter via the shared // helper so the `--all` path honors `frontmatter.embed_skip` the same @@ -466,7 +658,10 @@ async function embedAll( // (3000+/min for tier 1 = 50+/sec, 20 parallel is safely below) and // avoids overwhelming postgres connection pools. Users can tune via // GBRAIN_EMBED_CONCURRENCY env var based on their tier/infra. - const CONCURRENCY = parseInt(process.env.GBRAIN_EMBED_CONCURRENCY || '20', 10); + // Paced runs lower this to the resolved cap (the real lever vs pooler-slot + // starvation); unpaced keeps the env/default 20. + const CONCURRENCY = staleOpts?.paceMaxConcurrency + ?? parseInt(process.env.GBRAIN_EMBED_CONCURRENCY || '20', 10); async function embedOnePage(page: typeof pages[number]) { // #1737: bail before doing any work for this page if the run was aborted. @@ -475,7 +670,7 @@ async function embedAll( // target the correct (source_id, slug) row, not the 'default' source. const pageSourceId = page.source_id; const pageOpts = pageSourceId ? { sourceId: pageSourceId } : undefined; - const chunks = await engine.getChunks(page.slug, pageOpts); + const chunks = await observed(pacer, () => engine.getChunks(page.slug, pageOpts)); const toEmbed = chunks; // staleOnly path handled above via embedAllStale result.total_chunks += chunks.length; @@ -511,10 +706,12 @@ async function embedAll( embedding: embeddingMap.get(c.chunk_index) ?? undefined, token_count: c.token_count || Math.ceil(c.chunk_text.length / 4), })); - await engine.upsertChunks(page.slug, updated, pageOpts); + await observed(pacer, () => engine.upsertChunks(page.slug, updated, pageOpts)); // v0.41.31: stamp embedding provenance so a later model swap is // detectable as stale. - await engine.setPageEmbeddingSignature(page.slug, { sourceId: pageSourceId, signature }); + await observed(pacer, () => + engine.setPageEmbeddingSignature(page.slug, { sourceId: pageSourceId, signature }), + ); result.embedded += toEmbed.length; } catch (e: unknown) { serr(`\n Error embedding ${page.slug}: ${e instanceof Error ? e.message : e}`); @@ -523,6 +720,12 @@ async function embedAll( processed++; result.pages_processed++; onProgress?.(processed, pages.length, result.embedded); + // Cooperative DB-contention pace between pages (no-op when unpaced). + try { + await pacer.pace(signal); + } catch (e) { + if (!(e instanceof AbortError)) throw e; + } } // v0.41.15.0: sliding worker pool extracted into src/core/worker-pool.ts. @@ -576,6 +779,10 @@ async function embedAllStale( batchSize?: number; priority?: 'recent'; catchUp?: boolean; + /** DB-contention pacer (paced-backfill); no-op when pacing is off. */ + pacer?: DbPacer; + /** Resolved concurrency cap (E-1: the worker count, no separate permit). */ + paceMaxConcurrency?: number; }, signature?: string, externalSignal?: AbortSignal, @@ -626,7 +833,11 @@ async function embedAllStale( // (page_id, chunk_index). Each query finishes in <1s. // v0.41.18.0 (A13): --batch-size N CLI flag overrides hardcoded 2000 default. const PAGE_SIZE = staleOpts?.batchSize ?? 2000; - const CONCURRENCY = parseInt(process.env.GBRAIN_EMBED_CONCURRENCY || '20', 10); + // Paced runs lower concurrency to the resolved cap (E-1: worker count IS the + // lever on this single pool, no separate permit). Unpaced keeps env/default. + const CONCURRENCY = staleOpts?.paceMaxConcurrency + ?? parseInt(process.env.GBRAIN_EMBED_CONCURRENCY || '20', 10); + const pacer = staleOpts?.pacer ?? createNoopPacer(); // D3 + D3a + D8: wall-clock budget. 30 min default; env override. // #1946: --catch-up removes the wall-clock cap. The prior code set BUDGET_MS = @@ -640,9 +851,22 @@ async function embedAllStale( ? null : parseInt(process.env.GBRAIN_EMBED_TIME_BUDGET_MS || `${30 * 60 * 1000}`, 10); const budgetController = new AbortController(); - const budgetTimer = BUDGET_MS != null + const budgetStart = Date.now(); + let budgetTimer = BUDGET_MS != null ? setTimeout(() => budgetController.abort(), BUDGET_MS) : undefined; + // E-4 (paced-backfill): the budget measures WORK, not waiting. After each + // batch, re-arm the timer to fire at start + BUDGET + total-paced-sleep, so a + // contended DB that spends time in pace() sleeps converges instead of exiting + // having embedded little. No-op when unpaced (totalSleptMs stays 0) or in + // catch-up (no budget timer). + const rearmBudgetForPacing = (): void => { + if (BUDGET_MS == null) return; + const slept = pacer.snapshot().totalSleptMs; + if (budgetTimer) clearTimeout(budgetTimer); + const fireInMs = budgetStart + BUDGET_MS + slept - Date.now(); + budgetTimer = setTimeout(() => budgetController.abort(), Math.max(0, fireInMs)); + }; const budgetSignal = budgetController.signal; // #1737: the effective signal fires when EITHER the internal wall-clock // budget OR the caller's abort (worker timeout / lock loss / SIGTERM) fires. @@ -670,6 +894,33 @@ async function embedAllStale( // surfaces that loudly instead of looking like a clean run. let embedFailures = 0; + // E-3 (paced-backfill): bounded end-of-run re-entry. A longer paced run gives + // a live writer (sync / put_page) more time to insert NEW stale rows BEHIND + // the keyset cursor (TODOS:2301). When the cursor exhausts, re-scan from the + // start — capped at MAX_REENTRIES AND requiring forward progress (a pass that + // embeds 0 while count>0 stops) so a writer outrunning embed can't spin + // forever. + const MAX_REENTRIES = 3; + let reentries = 0; + let lastReentryEmbedded = 0; + const maybeReenter = async (): Promise => { + // Scoped to PACED runs: pacing lengthens the run, which is what widens the + // behind-cursor window. Unpaced runs keep prior (single-pass) behavior. + if (!pacer.snapshot().enabled) return false; + if (effectiveSignal.aborted) return false; + if (reentries >= MAX_REENTRIES) return false; + const remaining = await engine.countStaleChunks(sourceOpt); + if (remaining === 0) return false; + if (result.embedded === lastReentryEmbedded) return false; // no forward progress + lastReentryEmbedded = result.embedded; + reentries++; + afterPageId = 0; + afterChunkIndex = -1; + afterUpdatedAt = null; + serr(`\n [embed] re-entry ${reentries}/${MAX_REENTRIES}: ${remaining} stale chunk(s) appeared during the run; rescanning from start.`); + return true; + }; + try { // eslint-disable-next-line no-constant-condition while (true) { @@ -684,17 +935,22 @@ async function embedAllStale( break; } - const batch = await engine.listStaleChunks({ - batchSize: PAGE_SIZE, - afterPageId, - afterChunkIndex, - ...(orderBy === 'updated_desc' && { - orderBy, - afterUpdatedAt, + const batch = await observed(pacer, () => + engine.listStaleChunks({ + batchSize: PAGE_SIZE, + afterPageId, + afterChunkIndex, + ...(orderBy === 'updated_desc' && { + orderBy, + afterUpdatedAt, + }), + ...(sourceId && { sourceId }), }), - ...(sourceId && { sourceId }), - }); - if (batch.length === 0) break; + ); + if (batch.length === 0) { + if (await maybeReenter()) continue; + break; + } totalChunksLoaded += batch.length; // Advance cursor to last row in this batch. @@ -729,7 +985,7 @@ async function embedAllStale( try { const embeddings = await embedBatchWithBackoff(stale.map(c => c.chunk_text), { abortSignal: effectiveSignal }); // Re-fetch existing chunks and merge to avoid deleting non-stale chunks. - const existing = await engine.getChunks(slug, { sourceId: keySourceId }); + const existing = await observed(pacer, () => engine.getChunks(slug, { sourceId: keySourceId })); const staleIdxToEmbedding = new Map(); for (let j = 0; j < stale.length; j++) { staleIdxToEmbedding.set(stale[j].chunk_index, embeddings[j]); @@ -741,14 +997,16 @@ async function embedAllStale( embedding: staleIdxToEmbedding.get(c.chunk_index) ?? undefined, token_count: c.token_count || Math.ceil(c.chunk_text.length / 4), })); - await engine.upsertChunks(slug, merged, { sourceId: keySourceId }); + await observed(pacer, () => engine.upsertChunks(slug, merged, { sourceId: keySourceId })); // v0.41.31: stamp provenance after the page's chunks are embedded — // but only when EVERY chunk was stale (fully re-embedded this pass). // A partially-stale page keeps preserved chunks of unknown/old // provenance, so don't claim it's current. (After invalidate, a // signature-drifted page IS fully stale → this stamps it.) if (signature && stale.length === existing.length) { - await engine.setPageEmbeddingSignature(slug, { sourceId: keySourceId, signature }); + await observed(pacer, () => + engine.setPageEmbeddingSignature(slug, { sourceId: keySourceId, signature }), + ); } result.embedded += stale.length; } catch (e: unknown) { @@ -763,6 +1021,14 @@ async function embedAllStale( // Use staleCount as the estimated total for progress (not exact after // pagination starts, but directionally correct). onProgress?.(totalProcessedPages, Math.ceil(staleCount / PAGE_SIZE) * keys.length, result.embedded); + // Cooperative DB-contention pace between keys (no-op when unpaced). + // pace() throws AbortError on cancel; the pool stops claiming on the + // shared effectiveSignal, so swallow it here. + try { + await pacer.pace(effectiveSignal); + } catch (e) { + if (!(e instanceof AbortError)) throw e; + } } // v0.41.15.0: migrated to shared runSlidingPool. The pool checks @@ -778,8 +1044,14 @@ async function embedAllStale( failureLabel: (key) => key, }); + // E-4: extend the work budget by any paced-sleep time accrued this batch. + rearmBudgetForPacing(); + // If we got fewer rows than PAGE_SIZE, we've reached the end. - if (batch.length < PAGE_SIZE) break; + if (batch.length < PAGE_SIZE) { + if (await maybeReenter()) continue; + break; + } } } finally { if (budgetTimer) clearTimeout(budgetTimer); diff --git a/src/commands/jobs.ts b/src/commands/jobs.ts index a2d7db0a6..ec336f11b 100644 --- a/src/commands/jobs.ts +++ b/src/commands/jobs.ts @@ -8,6 +8,7 @@ import { MinionQueue } from '../core/minions/queue.ts'; import { MinionWorker } from '../core/minions/worker.ts'; import { WORKER_EXIT_RSS_WATCHDOG } from '../core/minions/worker-exit-codes.ts'; import type { MinionJob, MinionJobStatus } from '../core/minions/types.ts'; +import type { PaceKeyOverrides } from '../core/pace-mode.ts'; import { loadConfig, isThinClient } from '../core/config.ts'; import { callRemoteTool, unpackToolResult } from '../core/mcp-client.ts'; import { parseNiceValue, applyNiceness, getEffectiveNiceness, formatNice } from '../core/minions/niceness.ts'; @@ -1398,6 +1399,13 @@ export async function registerBuiltinHandlers( slugs: Array.isArray(job.data.slugs) ? (job.data.slugs as string[]) : undefined, all: !!job.data.all, stale: job.data.all ? false : (job.data.stale !== false), + sourceId: typeof job.data.sourceId === 'string' ? job.data.sourceId : undefined, + // CX1+CX5: pace overrides ride in the job payload as explicit overrides + // only; runEmbedCore re-resolves env > config > bundle at execution so + // GBRAIN_PACE_* still wins during an incident. + ...(job.data.pace && typeof job.data.pace === 'object' + ? { pace: job.data.pace as { perCallMode?: string; perCall?: PaceKeyOverrides } } + : {}), onProgress: (done, total, embedded) => { // Fire-and-forget: progress updates are best-effort and must not // block the worker loop. diff --git a/src/core/embed-stale.ts b/src/core/embed-stale.ts index 5ca95cd78..484f3ea55 100644 --- a/src/core/embed-stale.ts +++ b/src/core/embed-stale.ts @@ -20,6 +20,8 @@ import type { BrainEngine } from './engine.ts'; import type { ChunkInput } from './types.ts'; import { embedBatchWithBackoff } from '../commands/embed.ts'; +import { type DbPacer, createNoopPacer, observed } from './db-pacer.ts'; +import { AbortError } from './abort-check.ts'; /** Last visited (page_id, chunk_index) for keyset-resume across runs. */ export interface StaleCursor { @@ -59,6 +61,15 @@ export interface EmbedStaleOpts { * Omit to keep the legacy `embedding IS NULL`-only behavior. */ embeddingSignature?: string; + /** + * DB-contention pacer (paced-backfill). When enabled it (a) supplies the + * worker count via the caller passing `concurrency = bundle.maxConcurrency` + * — E-1: no separate permit on this single-pool path — and (b) the loop + * `observe()`s its DB-op latency and `pace()`s between keys. Omit (or pass a + * disabled bundle) for a no-op. The pacer is NEVER used to acquire permits + * here; concurrency is the worker count. + */ + pacer?: DbPacer; } export interface EmbedStaleResult { @@ -105,6 +116,9 @@ export async function embedStaleForSource( const signal = opts.signal; const embedFn = opts.embedFn ?? ((texts, fnOpts) => embedBatchWithBackoff(texts, { abortSignal: fnOpts.abortSignal })); + // Defaulted no-op when pacing is off, so the observe()/pace() call sites + // below are unconditional and cost ~nothing on the unpaced path. + const pacer = opts.pacer ?? createNoopPacer(); let afterPageId = opts.cursor?.afterPageId ?? 0; let afterChunkIndex = opts.cursor?.afterChunkIndex ?? -1; @@ -136,12 +150,14 @@ export async function embedStaleForSource( return result; } - const batch = await engine.listStaleChunks({ - batchSize, - afterPageId, - afterChunkIndex, - sourceId, - }); + const batch = await observed(pacer, () => + engine.listStaleChunks({ + batchSize, + afterPageId, + afterChunkIndex, + sourceId, + }), + ); if (batch.length === 0) { result.done = true; return result; @@ -177,7 +193,9 @@ export async function embedStaleForSource( stale.map((c) => c.chunk_text), { abortSignal: signal }, ); - const existing = await engine.getChunks(slug, { sourceId: keySourceId }); + const existing = await observed(pacer, () => + engine.getChunks(slug, { sourceId: keySourceId }), + ); const staleIdxToEmbedding = new Map(); for (let j = 0; j < stale.length; j++) { staleIdxToEmbedding.set(stale[j].chunk_index, embeddings[j]); @@ -189,13 +207,15 @@ export async function embedStaleForSource( embedding: staleIdxToEmbedding.get(c.chunk_index) ?? undefined, token_count: c.token_count || Math.ceil(c.chunk_text.length / 4), })); - await engine.upsertChunks(slug, merged, { sourceId: keySourceId }); + await observed(pacer, () => engine.upsertChunks(slug, merged, { sourceId: keySourceId })); // v0.41.31: stamp provenance only when EVERY chunk was stale (fully // re-embedded this pass) — a partially-stale page keeps preserved // chunks of unknown provenance, so don't claim current. After the // invalidate pass above, signature-drifted pages ARE fully stale. if (signature && stale.length === existing.length) { - await engine.setPageEmbeddingSignature(slug, { sourceId: keySourceId, signature }); + await observed(pacer, () => + engine.setPageEmbeddingSignature(slug, { sourceId: keySourceId, signature }), + ); } result.embedded += stale.length; result.pagesProcessed += 1; @@ -215,6 +235,15 @@ export async function embedStaleForSource( while (nextIdx < keys.length && !signal?.aborted) { const idx = nextIdx++; await embedOneKey(keys[idx]); + // Cooperative DB-contention pace between keys (no-op when unpaced). + // pace() throws AbortError on cancel — treat as graceful worker exit; + // the for(;;) loop sees signal.aborted and returns aborted next tick. + try { + await pacer.pace(signal); + } catch (e) { + if (e instanceof AbortError) return; + throw e; + } } } diff --git a/src/core/minions/handlers/embed-backfill.ts b/src/core/minions/handlers/embed-backfill.ts index 221d5b676..362ea470f 100644 --- a/src/core/minions/handlers/embed-backfill.ts +++ b/src/core/minions/handlers/embed-backfill.ts @@ -36,11 +36,14 @@ import { BudgetTracker, BudgetExhausted } from '../../budget/budget-tracker.ts'; import { withBudgetTracker } from '../../ai/gateway.ts'; import { embedStaleForSource } from '../../embed-stale.ts'; import { currentEmbeddingSignature } from '../../embedding.ts'; +import { type DbPacer, createDbPacer, createNoopPacer } from '../../db-pacer.ts'; +import { resolvePaceMode, loadPaceModeConfig, readPaceEnv } from '../../pace-mode.ts'; import type { BrainEngine } from '../../engine.ts'; import type { MinionJobContext } from '../types.ts'; +import { embedBackfillLockId, EMBED_BACKFILL_LOCK_TTL_MIN } from '../../embed-backfill-lock.ts'; + const DEFAULT_MAX_USD_PER_JOB = 10; -const EMBED_BACKFILL_LOCK_TTL_MIN = 60; export interface EmbedBackfillJobData { sourceId: string; @@ -61,9 +64,35 @@ export interface EmbedBackfillResult { budgetCapUsd?: number; } -/** Compose the lock id for embed-backfill, namespaced like sync's. */ -function embedBackfillLockId(sourceId: string): string { - return `gbrain-embed-backfill:${sourceId}`; +/** + * Resolve a DB-contention pacer from env > config > bundle for the prod + * backfill path. Fail-open: any error → no-op pacer (pacing never breaks a + * backfill). Returns the pacer + the resolved concurrency cap (E-1: the + * worker count for embedStaleForSource's single pool, no separate permit). + */ +async function resolveBackfillPacer( + engine: BrainEngine, + jobData: Record, +): Promise<{ pacer: DbPacer; concurrency?: number }> { + try { + const cfg = await loadPaceModeConfig(engine); + const { envMode, envOverrides } = readPaceEnv(); + const jobPace = (jobData.pace && typeof jobData.pace === 'object' + ? jobData.pace + : {}) as { perCallMode?: string; perCall?: Record }; + const knobs = resolvePaceMode({ + mode: cfg.mode, + configOverrides: cfg.configOverrides, + envMode, + envOverrides, + perCallMode: jobPace.perCallMode, + perCall: jobPace.perCall, + }); + if (!knobs.enabled) return { pacer: createNoopPacer() }; + return { pacer: createDbPacer({ bundle: knobs }), concurrency: knobs.maxConcurrency }; + } catch { + return { pacer: createNoopPacer() }; + } } /** Read embed.backfill_max_usd config or default. */ @@ -119,11 +148,18 @@ export function makeEmbedBackfillHandler(engine: BrainEngine) { label: `embed-backfill:${sourceId}`, }); + // paced-backfill: resolve env > config > bundle (env = incident escape + // hatch). No-op when off. This is the prod path that originally starved + // the supervisor, so pacing it is the headline win. + const { pacer, concurrency } = await resolveBackfillPacer(engine, job.data); + try { const result = await withBudgetTracker(tracker, async () => embedStaleForSource(engine, sourceId, { batchSize, signal: job.signal, + pacer, + ...(concurrency !== undefined && { concurrency }), // v0.41.31: re-embed pages whose model signature drifted + stamp // provenance as chunks land. embeddingSignature: currentEmbeddingSignature(), @@ -174,6 +210,7 @@ export function makeEmbedBackfillHandler(engine: BrainEngine) { } throw err; } finally { + pacer.dispose(); // ALWAYS release. Aborts, throws, budget-exhaust — all paths unwind here. try { await lock.release();