From 7968f8407709d42696efa60ecf0686b44dc5d4f5 Mon Sep 17 00:00:00 2001 From: Garry Tan Date: Wed, 17 Jun 2026 06:58:07 -0700 Subject: [PATCH] v0.42.49.0 feat(pace): native DB-contention pacing for embed/sync backfills (#2240) * feat(pace): composable DB-contention pacer primitive createDbPacer/createNoopPacer (concurrency permit + in-band EWMA + jittered cooperative sleep, abort-throws, fail-open) + named pace-mode bundles (env>config>bundle, default off) + shared embed-backfill lock key. * 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. * feat(sync): shared DB-pacer permit across parallel worker engines One pacer spans the per-worker PostgresEngines (the multi-pool permit case); observe() import writes, pace() between files, dispose on all exit paths. * docs(pace): CLAUDE.md Pace Mode section + regenerated llms bundle * fix(pace): pre-landing review fixes (Codex P1/P2) - never unref() the cooperative-sleep timer (could exit mid-sleep) - pace() excludes the wall-clock budget + re-arm after each sleep - pacing only lowers concurrency, never raises above an operator cap - serialized job pace resolves at config tier so GBRAIN_PACE_* still wins - --pace-max-concurrency consumes its value token * chore: bump version and changelog (v0.42.48.0) Native DB-contention pacing for embed/sync backfills. Co-Authored-By: Claude Opus 4.8 (1M context) * chore: bump version to v0.42.49.0 Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 17 + CLAUDE.md | 51 +++ TODOS.md | 30 ++ VERSION | 2 +- llms-full.txt | 51 +++ package.json | 2 +- src/commands/embed.ts | 351 ++++++++++++++++++-- src/commands/jobs.ts | 13 + src/commands/sync.ts | 53 ++- src/core/db-pacer.ts | 305 +++++++++++++++++ src/core/embed-backfill-lock.ts | 17 + src/core/embed-stale.ts | 47 ++- src/core/minions/handlers/embed-backfill.ts | 46 ++- src/core/pace-mode.ts | 286 ++++++++++++++++ test/db-pacer.test.ts | 202 +++++++++++ test/pace-mode.test.ts | 120 +++++++ 16 files changed, 1550 insertions(+), 43 deletions(-) create mode 100644 src/core/db-pacer.ts create mode 100644 src/core/embed-backfill-lock.ts create mode 100644 src/core/pace-mode.ts create mode 100644 test/db-pacer.test.ts create mode 100644 test/pace-mode.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index dba9421b8..af8687491 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,23 @@ All notable changes to GBrain will be documented in this file. +## [0.42.49.0] - 2026-06-16 + +**Big embed backfills and syncs now throttle themselves when the database gets busy, so clearing a backlog can't starve the job queue — no more external babysitter scripts.** A naive `gbrain embed --stale` or large `gbrain sync` against a PgBouncer transaction-mode pooler could saturate it and starve the minion supervisor's lock renewals, cascading `lock-renewal-failed` into dead jobs. The field workaround was an external wrapper that SIGSTOP/SIGCONT'd the process off a side-pool latency probe. That approach was blind (the side pool read low latency while the pool that mattered starved), unsafe (SIGSTOP can freeze a process mid-transaction holding locks), and couldn't touch peak pressure. gbrain now does this natively, and better. + +Pacing is **opt-in** (default `off`) and built on one composable primitive: it caps simultaneous in-flight DB writes (the real lever against pooler-slot starvation), measures the work's own query latency in-band (so it can never be blind), and sleeps cooperatively between safe points (never mid-transaction, so the lock heartbeat keeps firing). Turn it on per-run with `gbrain embed --stale --pace`, or set `pace.mode` in config to pace every embed path plus the production embed-backfill job automatically. `GBRAIN_PACE_*` env vars override config as an incident escape hatch. + +### Added +- **`--pace[=mode]` for `gbrain embed`** — `off`/`gentle`/`balanced`/`aggressive` bundles (bare `--pace` = balanced), plus `--pace-max-concurrency=N`. `--background` carries the explicit override into the queued `embed` job; the handler re-resolves env > config > bundle at execution. +- **`pace.mode` config + `GBRAIN_PACE_*` env** — config paces every `runEmbedCore` caller (cycle embed, catch-up, sync-auto-embed) and the prod `embed-backfill` job automatically; env beats config for incident response. +- **Composable `db-pacer` primitive** (`src/core/db-pacer.ts`) + named bundles (`src/core/pace-mode.ts`) — concurrency permit + in-band EWMA + jittered cooperative sleep, abort-throwing, fail-open. `sync` uses the shared permit across its parallel worker engines. +- **Pacing telemetry** — `EmbedResult.pacing` (cap, samples, EWMA latency, slept ms) in `--json`, plus a one-line stderr summary. + +### Changed +- **`gbrain embed --stale` now single-flights per source** using the same lock the `embed-backfill` job holds, so a hand-run backfill and a queued job can't grind the same source concurrently. Paced runs add a bounded end-of-run rescan (catches rows that landed behind the cursor during a longer run), and the embed time budget excludes paced-sleep time so a contended DB still converges instead of exiting early. + +### To take advantage of v0.42.49.0 +`gbrain upgrade`. Pacing is off by default — nothing changes until you opt in. To clear a big embed backlog safely on a busy pooler: `gbrain embed --stale --pace` (or `--pace=gentle` to be extra conservative). To pace the background embed-backfill job and every embed path automatically: `gbrain config set pace.mode balanced`. During an incident you can override without a redeploy: `GBRAIN_PACE_MODE=gentle` or `GBRAIN_PACE_MAX_CONCURRENCY=4`. ## [0.42.48.0] - 2026-06-16 **Brain repos harden themselves for durability the moment gbrain is given a PAT and a GitHub URL.** Fresh agents kept drifting out of sync with their knowledge-wiki git repos: writes sat local-only and never pushed, long-lived sessions edited a stale tree, and scratch output landed outside the repo and vanished. Now `gbrain sources add --url --pat-file

` auto-hardens the managed clone, and `gbrain sources harden ` runs the same audit idempotently against any source. Hardening is six always-on guarantees: it pulls current state (divergence-safe rebase that skips a dirty tree and never leaves a half-rebase), installs a local auto-push safety net, ships a committed `scripts/brain-commit-push.sh` that refuses to report success without a confirmed push, writes always-on durability rules into the agent's context file (deterministic filing from the canonical taxonomy, commit-and-push-never-deferred, pull-before-each-write-batch), registers a 30-minute background pull so an idle session can't go stale, and verifies push access up front. diff --git a/CLAUDE.md b/CLAUDE.md index e27e4b981..30c36e6df 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -401,6 +401,57 @@ hatches — no config-dashboard surface by design): | `GBRAIN_SYNC_YIELD_EVERY` | 64 | Yield the event loop (`setTimeout(0)`, NOT `setImmediate` — Bun starves the timers phase under a tight setImmediate loop) every N files so the lock-refresh `setInterval` heartbeat fires mid-import. | | `GBRAIN_LOCK_STEAL_GRACE_SECONDS` | derived (~600 at 30min TTL) | A holder that refreshed within this window is NOT stolen even if its TTL lapsed (starved-but-alive). Dead holders stop refreshing, age past the grace, and become stealable; TTL stays the backstop. | +## Pace Mode (DB-contention-aware backfill pacing) + +A naive `gbrain embed --stale` / large `sync` can saturate a PgBouncer +transaction-mode pooler and starve the minion supervisor's lock renewals +(`lock-renewal-failed` → dead jobs). Pacing is the native, composable fix — it +replaces external SIGSTOP/SIGCONT wrapper scripts. **Opt-in: default mode `off`.** + +The composable primitive is `src/core/db-pacer.ts` (`createDbPacer`): +- **Concurrency cap is the real lever** (caps simultaneous in-flight DB writes = + pooler slots held). Embed paths set their worker count to `maxConcurrency` + (single pool, no permit); `sync` uses the shared `acquire()` **permit** because + each parallel worker owns a separate engine (one budget must span pools). +- **In-band signal** (`observe(ms)` EWMA from the work's own queries — never + blind the way an out-of-band probe pool was). **No probe loop, no + `probeLatency` engine method.** +- **Cooperative `pace()` sleep** on `setTimeout` (keeps the lock heartbeat + firing), jittered to avoid a thundering-herd resume. `acquire()`/`pace()` throw + `AbortError` on cancel; everything else is fail-open (a pacer bug never kills a + backfill, never throws an unhandledRejection). + +Named bundles resolve through `src/core/pace-mode.ts` (`resolvePaceMode`), mirror +of the search-mode pattern but with **env ABOVE config** (incident escape hatch): + + per-call flag → GBRAIN_PACE_* env → config (pace.*) → PACE_BUNDLES[mode] → off + +| Knob | off | gentle | balanced | aggressive | +|---|---|---|---|---| +| `maxConcurrency` | (off) | 4 | 8 | 16 | +| `paceAtMs` (EWMA → sleep) | — | 250 | 500 | 1000 | +| `maxSleepMs` (jittered cap) | — | 2000 | 1500 | 1000 | + +**Surfaces.** `gbrain embed --stale --pace[=mode]` (bare `--pace` = balanced), +`--pace-max-concurrency=N`. `--background` carries explicit pace OVERRIDES (not +the resolved bundle) into the `embed` job payload; the handler re-resolves +env>config>bundle at execution so `GBRAIN_PACE_*` still wins (CX5). Config-level +`pace.mode` paces EVERY `runEmbedCore` caller (cycle embed, embed-catch-up, +sync-auto-embed) and the prod `embed-backfill` job automatically. `sync` reads +env/config. PGLite / mode `off` → no-op pacer. + +**Correctness fixes pacing bundles** (longer paced runs widen these): CLI +`embed --stale` single-flights via the SAME per-source lock key as the +`embed-backfill` handler (`src/core/embed-backfill-lock.ts`; all-source runs lock +every source in sorted order) so a hand-run backfill and a queued job can't race +the NULL→non-NULL upsert (`TODOS:2299`); a **bounded** end-of-run keyset re-entry +(max 3 + forward-progress, paced runs only) catches rows inserted behind the +cursor (`TODOS:2301`); and the embed wall-clock budget timer is re-armed around +`pace()` sleeps so paced time doesn't burn the work budget. + +`EmbedResult.pacing` carries the end-of-run telemetry (cap, samples, EWMA, slept +ms, max waiters) for `--json`; a one-line summary prints to stderr. + ## Build `bun build --compile --outfile bin/gbrain src/cli.ts` diff --git a/TODOS.md b/TODOS.md index ffbfaab58..8321af0e8 100644 --- a/TODOS.md +++ b/TODOS.md @@ -1,5 +1,35 @@ # TODOS +## Pace Mode follow-ups (filed v0.42.49.0) + +Deferred from the paced-backfill wave (CEO + eng review CLEARED). Core shipped: +`db-pacer` + `pace-mode` wired into embed (CLI + shared core + `embed-backfill` +job) and sync. See CLAUDE.md "Pace Mode". + +- [ ] **P2 — `doctor` pacing check (E2).** Detect a txn-mode pooler (port 6543) + running unpaced bulk and recommend `--pace`; optionally correlate recent + `minion_jobs` deaths with backfill windows. Where: `src/commands/doctor.ts`. +- [ ] **P2 — `--pace=auto` autotuned thresholds (E3).** Derive `paceAtMs`/cap from + observed baseline latency (rolling median) instead of fixed bundle values, + mirroring `gbrain search tune`. Needs a baseline window + cold-start default + + config persistence — not a small add. Where: `src/core/pace-mode.ts` + + `src/core/db-pacer.ts`. +- [ ] **P3 — First-class pacing in more minion job handlers (E5).** `embed-backfill` + is paced; extend to `extract`/`embed-catch-up`/contextual-reindex handlers with + supervisor-detection downgrade. Today these inherit config/env pacing only when + they call `runEmbedCore`. +- [ ] **P1-companion — Supervisor concurrency 3→2 + job-kind slot fairness (E7).** + The daemon-side root cause the external wrapper's probe was blind to: + `embed-backfill`/`autopilot-cycle` jobs can occupy all supervisor slots + (`:215` below). Pacing makes backfills safe; this fixes the residual death rate. + Where: `src/core/minions/supervisor.ts` + queue slot accounting. +- [ ] **P3 — `gbrain sync --pace` CLI flag.** Sync reads env/config pacing today; + add a per-run `--pace[=mode]` flag for symmetry with `embed`. Where: + `src/commands/sync.ts` arg parsing. +- [ ] **P3 — Real-PG e2e for pacing.** Gated on `DATABASE_URL`: paced + `embed --stale --pace --progress-json` caps concurrency + emits telemetry; + single-flight rejects a 2nd concurrent run; lock heartbeat advances during a + paced sleep (short-TTL). Unit coverage (`db-pacer`/`pace-mode`) already ships. ## brain-repo durability follow-ups (filed v0.42.48.0) - [ ] **P3 — gbrain write-path calls commit-push synchronously when durability is on.** diff --git a/VERSION b/VERSION index 7af6e6bde..bf317619b 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.42.48.0 +0.42.49.0 \ No newline at end of file diff --git a/llms-full.txt b/llms-full.txt index 879ef9ad0..2dcb09aad 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -550,6 +550,57 @@ hatches — no config-dashboard surface by design): | `GBRAIN_SYNC_YIELD_EVERY` | 64 | Yield the event loop (`setTimeout(0)`, NOT `setImmediate` — Bun starves the timers phase under a tight setImmediate loop) every N files so the lock-refresh `setInterval` heartbeat fires mid-import. | | `GBRAIN_LOCK_STEAL_GRACE_SECONDS` | derived (~600 at 30min TTL) | A holder that refreshed within this window is NOT stolen even if its TTL lapsed (starved-but-alive). Dead holders stop refreshing, age past the grace, and become stealable; TTL stays the backstop. | +## Pace Mode (DB-contention-aware backfill pacing) + +A naive `gbrain embed --stale` / large `sync` can saturate a PgBouncer +transaction-mode pooler and starve the minion supervisor's lock renewals +(`lock-renewal-failed` → dead jobs). Pacing is the native, composable fix — it +replaces external SIGSTOP/SIGCONT wrapper scripts. **Opt-in: default mode `off`.** + +The composable primitive is `src/core/db-pacer.ts` (`createDbPacer`): +- **Concurrency cap is the real lever** (caps simultaneous in-flight DB writes = + pooler slots held). Embed paths set their worker count to `maxConcurrency` + (single pool, no permit); `sync` uses the shared `acquire()` **permit** because + each parallel worker owns a separate engine (one budget must span pools). +- **In-band signal** (`observe(ms)` EWMA from the work's own queries — never + blind the way an out-of-band probe pool was). **No probe loop, no + `probeLatency` engine method.** +- **Cooperative `pace()` sleep** on `setTimeout` (keeps the lock heartbeat + firing), jittered to avoid a thundering-herd resume. `acquire()`/`pace()` throw + `AbortError` on cancel; everything else is fail-open (a pacer bug never kills a + backfill, never throws an unhandledRejection). + +Named bundles resolve through `src/core/pace-mode.ts` (`resolvePaceMode`), mirror +of the search-mode pattern but with **env ABOVE config** (incident escape hatch): + + per-call flag → GBRAIN_PACE_* env → config (pace.*) → PACE_BUNDLES[mode] → off + +| Knob | off | gentle | balanced | aggressive | +|---|---|---|---|---| +| `maxConcurrency` | (off) | 4 | 8 | 16 | +| `paceAtMs` (EWMA → sleep) | — | 250 | 500 | 1000 | +| `maxSleepMs` (jittered cap) | — | 2000 | 1500 | 1000 | + +**Surfaces.** `gbrain embed --stale --pace[=mode]` (bare `--pace` = balanced), +`--pace-max-concurrency=N`. `--background` carries explicit pace OVERRIDES (not +the resolved bundle) into the `embed` job payload; the handler re-resolves +env>config>bundle at execution so `GBRAIN_PACE_*` still wins (CX5). Config-level +`pace.mode` paces EVERY `runEmbedCore` caller (cycle embed, embed-catch-up, +sync-auto-embed) and the prod `embed-backfill` job automatically. `sync` reads +env/config. PGLite / mode `off` → no-op pacer. + +**Correctness fixes pacing bundles** (longer paced runs widen these): CLI +`embed --stale` single-flights via the SAME per-source lock key as the +`embed-backfill` handler (`src/core/embed-backfill-lock.ts`; all-source runs lock +every source in sorted order) so a hand-run backfill and a queued job can't race +the NULL→non-NULL upsert (`TODOS:2299`); a **bounded** end-of-run keyset re-entry +(max 3 + forward-progress, paced runs only) catches rows inserted behind the +cursor (`TODOS:2301`); and the embed wall-clock budget timer is re-armed around +`pace()` sleeps so paced time doesn't burn the work budget. + +`EmbedResult.pacing` carries the end-of-run telemetry (cap, samples, EWMA, slept +ms, max waiters) for `--json`; a one-line summary prints to stderr. + ## Build `bun build --compile --outfile bin/gbrain src/cli.ts` diff --git a/package.json b/package.json index b7d577b94..4c08fb119 100644 --- a/package.json +++ b/package.json @@ -143,5 +143,5 @@ "bun": ">=1.3.10" }, "license": "MIT", - "version": "0.42.48.0" + "version": "0.42.49.0" } diff --git a/src/commands/embed.ts b/src/commands/embed.ts index 3afefab03..6d8816a6e 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,33 @@ 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; + }; + /** + * When the pace overrides were SERIALIZED from a background-job payload (not + * typed at an interactive CLI), resolve them at the config tier so + * `GBRAIN_PACE_*` on the worker still overrides at execution (Codex P2). Set + * by the `embed` job handler; unset for interactive CLI runs. + */ + paceFromBackground?: boolean; + /** + * 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 +130,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 +261,118 @@ 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(); + // Codex P2: an interactive CLI flag (--pace) is the most immediate + // intent and sits at the per-call tier (beats env). But a flag + // SERIALIZED into a background job payload must sit at the CONFIG tier + // so GBRAIN_PACE_* on the worker can still override it at execution + // (incident escape hatch). paceFromBackground distinguishes the two. + const fromBg = !!opts.paceFromBackground; + const knobs = resolvePaceMode({ + mode: fromBg ? (opts.pace?.perCallMode ?? cfg.mode) : cfg.mode, + configOverrides: fromBg + ? { ...cfg.configOverrides, ...(opts.pace?.perCall ?? {}) } + : cfg.configOverrides, + envMode, + envOverrides, + perCallMode: fromBg ? undefined : opts.pace?.perCallMode, + perCall: fromBg ? undefined : 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 +382,39 @@ 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; + i++; // consume the value token so positional parsing can't read it as a slug (Codex P2) + } + } + 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 +433,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 +460,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 +616,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 +647,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 +674,13 @@ 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. Codex P2: only ever LOWER — + // never raise above an operator's existing env cap. + const BASE_CONCURRENCY = parseInt(process.env.GBRAIN_EMBED_CONCURRENCY || '20', 10); + const CONCURRENCY = staleOpts?.paceMaxConcurrency + ? Math.min(BASE_CONCURRENCY, staleOpts.paceMaxConcurrency) + : BASE_CONCURRENCY; async function embedOnePage(page: typeof pages[number]) { // #1737: bail before doing any work for this page if the run was aborted. @@ -475,7 +689,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 +725,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 +739,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 +798,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 +852,14 @@ 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). Codex P2: pacing only ever + // LOWERS concurrency — never raise above an operator's existing env cap. + const BASE_CONCURRENCY = parseInt(process.env.GBRAIN_EMBED_CONCURRENCY || '20', 10); + const CONCURRENCY = staleOpts?.paceMaxConcurrency + ? Math.min(BASE_CONCURRENCY, staleOpts.paceMaxConcurrency) + : BASE_CONCURRENCY; + 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 +873,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 +916,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 +957,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 +1007,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 +1019,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 +1043,17 @@ 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). + // E-4 (Codex P1): pace() is subject to the EXTERNAL abort only, NOT the + // wall-clock budget — a contended DB's sleep must not be cut by the + // budget timer before its time is credited. Re-arm the budget right + // after each sleep so accrued sleep never eats into work time. + try { + await pacer.pace(externalSignal); + rearmBudgetForPacing(); + } catch (e) { + if (!(e instanceof AbortError)) throw e; + } } // v0.41.15.0: migrated to shared runSlidingPool. The pool checks @@ -778,8 +1069,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..029b0c167 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,18 @@ 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 }, + // Serialized from the queued payload → config tier so GBRAIN_PACE_* + // on the worker still wins at execution (Codex P2 escape hatch). + paceFromBackground: true, + } + : {}), onProgress: (done, total, embedded) => { // Fire-and-forget: progress updates are best-effort and must not // block the worker loop. diff --git a/src/commands/sync.ts b/src/commands/sync.ts index 09f8bafa9..d788cd201 100644 --- a/src/commands/sync.ts +++ b/src/commands/sync.ts @@ -83,6 +83,9 @@ import { type OpCheckpointKey, } from '../core/op-checkpoint.ts'; import { registerCleanup } from '../core/process-cleanup.ts'; +import { type DbPacer, createDbPacer, createNoopPacer, observed } from '../core/db-pacer.ts'; +import { resolvePaceMode, loadPaceModeConfig, readPaceEnv } from '../core/pace-mode.ts'; +import { AbortError } from '../core/abort-check.ts'; /** * v0.42.x (#1794) -- resumable incremental sync checkpoint. @@ -2366,6 +2369,27 @@ async function performSyncInner(engine: BrainEngine, opts: SyncOpts): Promise config > bundle (env = incident escape + // hatch); fail-open so pacing never breaks a sync. + let pacer: DbPacer = createNoopPacer(); + try { + const pcfg = await loadPaceModeConfig(engine); + const { envMode, envOverrides } = readPaceEnv(); + const knobs = resolvePaceMode({ + mode: pcfg.mode, + configOverrides: pcfg.configOverrides, + envMode, + envOverrides, + }); + if (knobs.enabled) pacer = createDbPacer({ bundle: knobs }); + } catch { + pacer = createNoopPacer(); + } + async function importOnePath(eng: BrainEngine, path: string): Promise { const filePath = join(syncRepoPath, path); if (!existsSync(filePath)) { @@ -2396,13 +2420,24 @@ async function performSyncInner(engine: BrainEngine, opts: SyncOpts): Promise1 / --all the stuck file is the // begin-line with no matching completion in the in-flight set. if (process.env.GBRAIN_SYNC_TRACE) serr(`[sync] begin import: ${path}`); + // paced-backfill: acquire a DB-write permit (caps total concurrent writes + // across all worker engines). Throws AbortError on cancel while waiting — + // treat as a clean skip; the worker loop sees signal.aborted next tick. + let permit; + try { + permit = await pacer.acquire(opts.signal); + } catch (e) { + if (e instanceof AbortError) return; + throw e; + } try { // v0.18.0+ multi-source: thread `opts.sourceId` so per-page tx writes // (putPage / getTags / addTag / removeTag / deleteChunks / upsertChunks // / addLink) target (sourceId, slug). Pre-fix the schema DEFAULT // 'default' was applied even for non-default sources, fabricating // duplicate rows that crashed bare-slug subqueries with Postgres 21000. - const result = await importFile(eng, filePath, path, { noEmbed, sourceId: opts.sourceId, activePack: syncActivePack }); + const result = await observed(pacer, () => + importFile(eng, filePath, path, { noEmbed, sourceId: opts.sourceId, activePack: syncActivePack })); if (result.status === 'imported') { chunksCreated += result.chunks; pagesAffected.push(result.slug); @@ -2427,12 +2462,23 @@ async function performSyncInner(engine: BrainEngine, opts: SyncOpts): Promise; + /** Feed an observed DB-op latency sample (ms). NaN / negative is ignored. */ + observe(latencyMs: number): void; + /** + * Cooperative sleep when recent latency is high (EWMA > `paceAtMs`), jittered + * per call and capped at `maxSleepMs`. No-op when latency is fine. THROWS + * `AbortError` if `signal` fires during the sleep. + */ + pace(signal?: AbortSignal): Promise; + snapshot(): PaceSnapshot; + /** Stop accepting waiters and release any blocked acquirers. Idempotent. */ + dispose(): void; +} + +/** Test/impl seams. Production defaults read real timers + Math.random. */ +export interface DbPacerSeams { + /** Sleep `ms`, rejecting `AbortError` if `signal` fires first. */ + sleep?: (ms: number, signal?: AbortSignal) => Promise; + /** Jitter source in [0, 1). */ + rng?: () => number; +} + +export interface CreateDbPacerOpts extends DbPacerSeams { + bundle: PaceBundle; +} + +interface Waiter { + resolve: (p: Permit) => void; + reject: (e: unknown) => void; + onAbort?: () => void; + cleanup?: () => void; + signal?: AbortSignal; +} + +/** + * Default abortable sleep. Resolves after `ms`; rejects `AbortError` if the + * signal fires first. `ms <= 0` resolves on the next microtask. + */ +function defaultSleep(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted) return Promise.reject(new AbortError(abortReason(signal))); + if (!(ms > 0)) return Promise.resolve(); + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + cleanup(); + resolve(); + }, ms); + // NOTE: do NOT unref() this timer. The bulk loop is awaiting it, so if it + // were unref'd and no other referenced handle existed, the process could + // exit mid-sleep and truncate the backfill (Codex P1). + const onAbort = () => { + cleanup(); + reject(new AbortError(abortReason(signal))); + }; + const cleanup = () => { + clearTimeout(timer); + signal?.removeEventListener('abort', onAbort); + }; + signal?.addEventListener('abort', onAbort, { once: true }); + }); +} + +function abortReason(signal?: AbortSignal | null): string { + const r = signal?.reason; + if (r instanceof Error) return r.message; + return String(r ?? 'aborted'); +} + +const NOOP_PERMIT: Permit = { release() {} }; + +/** No-op pacer: unbounded concurrency, zero sleeps. For `off` mode / PGLite. */ +export function createNoopPacer(): DbPacer { + return { + acquire: () => Promise.resolve(NOOP_PERMIT), + observe: () => {}, + pace: () => Promise.resolve(), + snapshot: () => ({ + enabled: false, + maxConcurrency: 0, + active: 0, + ewmaMs: null, + totalSleptMs: 0, + sleepCount: 0, + maxWaiters: 0, + sampleCount: 0, + }), + dispose: () => {}, + }; +} + +export function createDbPacer(opts: CreateDbPacerOpts): DbPacer { + const { bundle } = opts; + if (!bundle.enabled) return createNoopPacer(); + + const max = Math.max(1, Math.floor(bundle.maxConcurrency)); + const paceAtMs = Math.max(0, bundle.paceAtMs); + const maxSleepMs = Math.max(0, bundle.maxSleepMs); + const alpha = bundle.ewmaAlpha > 0 && bundle.ewmaAlpha <= 1 ? bundle.ewmaAlpha : 0.3; + const sleep = opts.sleep ?? defaultSleep; + const rng = opts.rng ?? Math.random; + + let active = 0; + let disposed = false; + let ewma: number | null = null; + let totalSleptMs = 0; + let sleepCount = 0; + let maxWaiters = 0; + let sampleCount = 0; + const waiters: Waiter[] = []; + + function makePermit(): Permit { + let released = false; + return { + release() { + if (released) return; + released = true; + // Hand the slot to the next waiter without dropping `active`, so the + // cap is never momentarily exceeded. + const next = waiters.shift(); + if (next) { + next.cleanup?.(); + next.resolve(makePermit()); + } else { + active = Math.max(0, active - 1); + } + }, + }; + } + + async function acquire(signal?: AbortSignal): Promise { + try { + if (signal?.aborted) throw new AbortError(abortReason(signal)); + if (disposed) return NOOP_PERMIT; + if (active < max) { + active++; + return makePermit(); + } + return await new Promise((resolve, reject) => { + const waiter: Waiter = { resolve, reject, signal }; + const onAbort = () => { + const i = waiters.indexOf(waiter); + if (i >= 0) waiters.splice(i, 1); + reject(new AbortError(abortReason(signal))); + }; + waiter.onAbort = onAbort; + waiter.cleanup = () => signal?.removeEventListener('abort', onAbort); + signal?.addEventListener('abort', onAbort, { once: true }); + waiters.push(waiter); + if (waiters.length > maxWaiters) maxWaiters = waiters.length; + }); + } catch (err) { + // Abort must propagate so the loop can't fall into a DB call. + if (err instanceof AbortError) throw err; + // Anything else: fail-open (a pacer bug must not kill the backfill). + return NOOP_PERMIT; + } + } + + function observe(latencyMs: number): void { + try { + if (typeof latencyMs !== 'number' || !Number.isFinite(latencyMs) || latencyMs < 0) return; + ewma = ewma === null ? latencyMs : alpha * latencyMs + (1 - alpha) * ewma; + sampleCount++; + } catch { + /* fail-open */ + } + } + + async function pace(signal?: AbortSignal): Promise { + let ms = 0; + try { + if (ewma === null || ewma <= paceAtMs || maxSleepMs <= 0) return; + const base = Math.min(maxSleepMs, ewma); + // Jitter to 50-100% of base so N workers don't resume in lockstep. + const jitter = 0.5 + 0.5 * clamp01(rng()); + ms = Math.round(base * jitter); + if (ms <= 0) return; + } catch { + return; // fail-open: never let knob math kill the loop + } + try { + await sleep(ms, signal); + totalSleptMs += ms; + sleepCount++; + } catch (err) { + // Abort must propagate (a cancel can't be swallowed); any other sleep + // failure is fail-open — a pacer bug never kills the backfill. + if (err instanceof AbortError) throw err; + } + } + + function snapshot(): PaceSnapshot { + return { + enabled: true, + maxConcurrency: max, + active, + ewmaMs: ewma, + totalSleptMs, + sleepCount, + maxWaiters, + sampleCount, + }; + } + + function dispose(): void { + if (disposed) return; + disposed = true; + // Release blocked acquirers with no-op permits so their loops end cleanly + // rather than hanging forever. + while (waiters.length > 0) { + const w = waiters.shift(); + w?.cleanup?.(); + w?.resolve(NOOP_PERMIT); + } + } + + return { acquire, observe, pace, snapshot, dispose }; +} + +function clamp01(n: number): number { + if (!Number.isFinite(n)) return 0; + if (n < 0) return 0; + if (n > 1) return 1; + return n; +} + +/** + * Convenience: time a DB op and feed its latency to the pacer. Keeps call sites + * a one-liner — `await observed(pacer, () => engine.upsertChunks(...))`. + */ +export async function observed(pacer: DbPacer, fn: () => Promise): Promise { + const t0 = Date.now(); + try { + return await fn(); + } finally { + pacer.observe(Date.now() - t0); + } +} + +// anySignal is re-exported intentionally so call sites that compose a budget +// signal with the pacer can import from one place. +export { anySignal }; diff --git a/src/core/embed-backfill-lock.ts b/src/core/embed-backfill-lock.ts new file mode 100644 index 000000000..af9ceabc6 --- /dev/null +++ b/src/core/embed-backfill-lock.ts @@ -0,0 +1,17 @@ +/** + * Shared lock identity for embed backfills (paced-backfill E-2). + * + * The per-source lock key + TTL live here, in a zero-dependency module, so BOTH + * the `embed-backfill` minion handler AND the CLI `embed --stale` single-flight + * take the SAME key — a hand-run backfill and a queued job are then mutually + * exclusive per source. Kept dependency-free to avoid an import cycle between + * embed.ts, embed-stale.ts, and the handler. + */ + +/** Per-source embed-backfill lock id, namespaced like sync's. */ +export function embedBackfillLockId(sourceId: string): string { + return `gbrain-embed-backfill:${sourceId}`; +} + +/** Lock TTL (minutes) for embed backfills. */ +export const EMBED_BACKFILL_LOCK_TTL_MIN = 60; 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 8fb6d52ac..30ffd6ed5 100644 --- a/src/core/minions/handlers/embed-backfill.ts +++ b/src/core/minions/handlers/embed-backfill.ts @@ -36,12 +36,15 @@ 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 { parseUsdLimit, usdLimitToCap, resolveSpendPosture } from '../../spend-posture.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; @@ -62,9 +65,36 @@ 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 }; + // Codex P2: serialized job pace sits at the CONFIG tier (job choice beats + // standing config, but env beats both) so GBRAIN_PACE_* on the worker is a + // real incident escape hatch for an already-queued job. + const knobs = resolvePaceMode({ + mode: jobPace.perCallMode ?? cfg.mode, + configOverrides: { ...cfg.configOverrides, ...(jobPace.perCall ?? {}) }, + envMode, + envOverrides, + }); + if (!knobs.enabled) return { pacer: createNoopPacer() }; + return { pacer: createDbPacer({ bundle: knobs }), concurrency: knobs.maxConcurrency }; + } catch { + return { pacer: createNoopPacer() }; + } } /** @@ -129,11 +159,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(), @@ -184,6 +221,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(); diff --git a/src/core/pace-mode.ts b/src/core/pace-mode.ts new file mode 100644 index 000000000..7daf1da33 --- /dev/null +++ b/src/core/pace-mode.ts @@ -0,0 +1,286 @@ +/** + * DB-contention pacing mode bundles (paced-backfill internals). + * + * Named modes that bundle the DB-pacer knobs into a single config key so an + * operator picks once and stops thinking about it. Mirrors the v0.32.3 + * search-mode pattern at `src/core/search/mode.ts` (named bundle + per-key + * overrides + per-call opts), with ONE deliberate difference: the resolution + * chain puts ENV ABOVE CONFIG so `GBRAIN_PACE_*` is a real incident escape + * hatch — an operator can override bad production config without a redeploy. + * + * per-call flag → GBRAIN_PACE_* env → config (pace.*) → MODE_BUNDLES[mode] → off + * + * `DEFAULT_PACE_MODE` is `'off'`: pacing is strictly opt-in and never changes + * default behavior. `off` resolves to `enabled: false`, which the DB-pacer + * turns into a no-op (unbounded concurrency, zero sleeps). + * + * This module is PURE — no DB calls, no env reads inside `resolvePaceMode`. + * The caller pre-loads config (`loadPaceModeConfig`) and env (`readPaceEnv`) + * and passes them in, so the resolver is trivially testable. + */ + +export type PaceMode = 'off' | 'gentle' | 'balanced' | 'aggressive'; + +export const PACE_MODES: ReadonlyArray = Object.freeze([ + 'off', + 'gentle', + 'balanced', + 'aggressive', +]); + +export const DEFAULT_PACE_MODE: PaceMode = 'off'; + +/** + * A complete knob set for one pace mode. Every field is required so the bundle + * is self-contained and per-key overrides are obvious diffs. + */ +export interface PaceBundle { + /** Master switch. `false` ⇒ the DB-pacer is a no-op (off mode / PGLite). */ + enabled: boolean; + /** + * Primary lever: cap on simultaneous in-flight DB writes. For single-pool + * paths (embed) this becomes the worker count; for multi-pool paths (sync) + * it's enforced by the shared `acquire()` permit. The real defense against + * pooler-slot starvation. + */ + maxConcurrency: number; + /** EWMA of observed DB-op latency (ms) above which `pace()` starts sleeping. */ + paceAtMs: number; + /** Ceiling on a single cooperative sleep (ms). Jittered downward per call. */ + maxSleepMs: number; + /** EWMA smoothing factor in (0, 1]. Higher = reacts faster to recent latency. */ + ewmaAlpha: number; +} + +/** + * The four bundles. Frozen at import so a typo can't redefine "balanced" to + * mean different things on different installs. `off` is the disabled sentinel; + * its numeric fields are inert (the pacer short-circuits on `enabled: false`). + * + * Concurrency picks are deliberately well below the default embed concurrency + * (`GBRAIN_EMBED_CONCURRENCY`, 20) — the whole point is to hold fewer pooler + * slots than an unpaced run. + */ +export const PACE_BUNDLES: Readonly>> = Object.freeze({ + off: Object.freeze({ + enabled: false, + maxConcurrency: 0, + paceAtMs: 0, + maxSleepMs: 0, + ewmaAlpha: 0, + }), + gentle: Object.freeze({ + enabled: true, + maxConcurrency: 4, + paceAtMs: 250, + maxSleepMs: 2000, + ewmaAlpha: 0.3, + }), + balanced: Object.freeze({ + enabled: true, + maxConcurrency: 8, + paceAtMs: 500, + maxSleepMs: 1500, + ewmaAlpha: 0.3, + }), + aggressive: Object.freeze({ + enabled: true, + maxConcurrency: 16, + paceAtMs: 1000, + maxSleepMs: 1000, + ewmaAlpha: 0.3, + }), +}); + +export function isPaceMode(x: unknown): x is PaceMode { + return typeof x === 'string' && (PACE_MODES as ReadonlyArray).includes(x); +} + +/** + * Per-key overrides (from the config table OR from env). Every field optional; + * undefined ⇒ fall through to the next precedence layer. + */ +export interface PaceKeyOverrides { + enabled?: boolean; + maxConcurrency?: number; + paceAtMs?: number; + maxSleepMs?: number; + ewmaAlpha?: number; +} + +/** + * Resolve the active pace knob set. Pure: the caller supplies the resolved + * config + env layers. + * + * Mode precedence: perCallMode → envMode → mode(config) → DEFAULT_PACE_MODE + * Knob precedence: perCall → env → config → MODE_BUNDLES[mode] + * + * Env above config is intentional (incident escape hatch — see file header). + */ +export interface ResolvePaceModeInput { + /** `config.pace.mode`. */ + mode?: string; + /** `GBRAIN_PACE_MODE`. */ + envMode?: string; + /** `--pace=` (or bare `--pace` resolved to 'balanced' by the caller). */ + perCallMode?: string; + /** Per-key overrides from the config table. */ + configOverrides?: PaceKeyOverrides; + /** Per-key overrides from `GBRAIN_PACE_*` env. */ + envOverrides?: PaceKeyOverrides; + /** Per-call overrides (e.g. `--pace-max-concurrency`). */ + perCall?: PaceKeyOverrides; +} + +export interface ResolvedPaceKnobs extends PaceBundle { + /** Which bundle supplied the defaults (after fallback). */ + resolved_mode: PaceMode; + /** True if the resolved mode string was a recognized PaceMode. */ + mode_valid: boolean; +} + +export function resolvePaceMode(input: ResolvePaceModeInput): ResolvedPaceKnobs { + const rawMode = + firstString(input.perCallMode) ?? firstString(input.envMode) ?? firstString(input.mode); + const normalized = rawMode ? rawMode.trim().toLowerCase() : ''; + const valid = isPaceMode(normalized); + const resolved_mode: PaceMode = valid ? (normalized as PaceMode) : DEFAULT_PACE_MODE; + const bundle = PACE_BUNDLES[resolved_mode]; + + const pc = input.perCall ?? {}; + const env = input.envOverrides ?? {}; + const cfg = input.configOverrides ?? {}; + + const pick = (key: K): PaceBundle[K] => { + if (pc[key] !== undefined) return pc[key] as PaceBundle[K]; + if (env[key] !== undefined) return env[key] as PaceBundle[K]; + if (cfg[key] !== undefined) return cfg[key] as PaceBundle[K]; + return bundle[key]; + }; + + const enabled = pick('enabled'); + // Clamp to safe ranges so a fat-fingered override can't disable the cap + // (maxConcurrency must be >= 1 when enabled) or wedge the EWMA. + let maxConcurrency = pick('maxConcurrency'); + if (enabled && (!Number.isFinite(maxConcurrency) || maxConcurrency < 1)) { + maxConcurrency = bundle.enabled ? bundle.maxConcurrency : 8; + } + let ewmaAlpha = pick('ewmaAlpha'); + if (!Number.isFinite(ewmaAlpha) || ewmaAlpha <= 0 || ewmaAlpha > 1) { + ewmaAlpha = bundle.enabled ? bundle.ewmaAlpha : 0.3; + } + const paceAtMs = Math.max(0, pick('paceAtMs')); + const maxSleepMs = Math.max(0, pick('maxSleepMs')); + + return { + enabled, + maxConcurrency, + paceAtMs, + maxSleepMs, + ewmaAlpha, + resolved_mode, + mode_valid: valid, + }; +} + +function firstString(v: unknown): string | undefined { + return typeof v === 'string' && v.trim().length > 0 ? v : undefined; +} + +/** Config-table key for the mode selection (separate from the override keys). */ +export const PACE_MODE_KEY = 'pace.mode'; + +/** Per-knob config keys this module reads. Used by a future `--reset`. */ +export const PACE_MODE_CONFIG_KEYS: ReadonlyArray = Object.freeze([ + 'pace.enabled', + 'pace.max_concurrency', + 'pace.pace_at_ms', + 'pace.max_sleep_ms', + 'pace.ewma_alpha', +]); + +/** Build PaceKeyOverrides from a flat config-table snapshot (sparse). */ +export function loadOverridesFromConfig( + configMap: Record, +): PaceKeyOverrides { + return parseOverrides((k) => configMap[k], { + enabled: 'pace.enabled', + maxConcurrency: 'pace.max_concurrency', + paceAtMs: 'pace.pace_at_ms', + maxSleepMs: 'pace.max_sleep_ms', + ewmaAlpha: 'pace.ewma_alpha', + }); +} + +/** Build PaceKeyOverrides from `GBRAIN_PACE_*` env (incident escape hatch). */ +export function readPaceEnv(env: Record = process.env): { + envMode?: string; + envOverrides: PaceKeyOverrides; +} { + const overrides = parseOverrides((k) => env[k], { + enabled: 'GBRAIN_PACE_ENABLED', + maxConcurrency: 'GBRAIN_PACE_MAX_CONCURRENCY', + paceAtMs: 'GBRAIN_PACE_AT_MS', + maxSleepMs: 'GBRAIN_PACE_MAX_SLEEP_MS', + ewmaAlpha: 'GBRAIN_PACE_EWMA_ALPHA', + }); + return { envMode: env.GBRAIN_PACE_MODE, envOverrides: overrides }; +} + +function parseOverrides( + get: (k: string) => string | undefined, + keys: Record, +): PaceKeyOverrides { + const out: PaceKeyOverrides = {}; + const en = get(keys.enabled); + if (en !== undefined) out.enabled = en === '1' || en.toLowerCase() === 'true'; + const mc = get(keys.maxConcurrency); + if (mc !== undefined) { + const n = parseInt(mc, 10); + if (Number.isFinite(n) && n >= 1 && n <= 256) out.maxConcurrency = n; + } + const pa = get(keys.paceAtMs); + if (pa !== undefined) { + const n = parseInt(pa, 10); + if (Number.isFinite(n) && n >= 0) out.paceAtMs = n; + } + const ms = get(keys.maxSleepMs); + if (ms !== undefined) { + const n = parseInt(ms, 10); + if (Number.isFinite(n) && n >= 0) out.maxSleepMs = n; + } + const ea = get(keys.ewmaAlpha); + if (ea !== undefined) { + const n = parseFloat(ea); + if (Number.isFinite(n) && n > 0 && n <= 1) out.ewmaAlpha = n; + } + return out; +} + +/** + * Load the live pace config (mode + per-key overrides) from the brain engine. + * Errors swallowed → mode-bundle defaults (the config table may predate this + * feature on old brains). Does NOT read env — the caller layers `readPaceEnv` + * on top so env can beat config. + */ +export async function loadPaceModeConfig(engine: { + getConfig(key: string): Promise; +}): Promise<{ mode?: string; configOverrides: PaceKeyOverrides }> { + const safeGet = async (k: string): Promise => { + try { + const v = await engine.getConfig(k); + return typeof v === 'string' ? v : undefined; + } catch { + return undefined; + } + }; + const [mode, ...vals] = await Promise.all([ + safeGet(PACE_MODE_KEY), + ...PACE_MODE_CONFIG_KEYS.map(safeGet), + ]); + const configMap: Record = {}; + PACE_MODE_CONFIG_KEYS.forEach((key, i) => { + if (vals[i] !== undefined) configMap[key] = vals[i]; + }); + return { mode, configOverrides: loadOverridesFromConfig(configMap) }; +} diff --git a/test/db-pacer.test.ts b/test/db-pacer.test.ts new file mode 100644 index 000000000..c5168d52a --- /dev/null +++ b/test/db-pacer.test.ts @@ -0,0 +1,202 @@ +/** + * Pins the DB-pacer primitive: concurrency cap (the real lever), abort-throws, + * in-band EWMA, jittered cooperative sleep, fail-open, and the no-op path. + * + * Uses a fake sleep + deterministic rng so there's no real wall-clock wait and + * no flakiness — the whole suite runs in microtask time. + */ +import { describe, expect, test } from 'bun:test'; +import { createDbPacer, createNoopPacer, type DbPacer } from '../src/core/db-pacer.ts'; +import { AbortError } from '../src/core/abort-check.ts'; +import { PACE_BUNDLES, type PaceBundle } from '../src/core/pace-mode.ts'; + +const BUNDLE: PaceBundle = { + enabled: true, + maxConcurrency: 2, + paceAtMs: 100, + maxSleepMs: 1000, + ewmaAlpha: 0.5, +}; + +/** A fake sleep that resolves immediately but still honors abort. */ +function fakeSleep(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted) return Promise.reject(new AbortError('aborted')); + return Promise.resolve(); +} + +function pacer(over: Partial = {}, rng = () => 0.5): DbPacer { + return createDbPacer({ bundle: { ...BUNDLE, ...over }, sleep: fakeSleep, rng }); +} + +describe('concurrency cap (the real lever)', () => { + test('caps simultaneous permits at maxConcurrency; release hands off FIFO', async () => { + const p = pacer({ maxConcurrency: 2 }); + const a = await p.acquire(); + const b = await p.acquire(); + expect(p.snapshot().active).toBe(2); + + let cResolved = false; + const cPromise = p.acquire().then((permit) => { + cResolved = true; + return permit; + }); + // c is blocked: still 2 active, 1 waiter. + await Promise.resolve(); + expect(cResolved).toBe(false); + expect(p.snapshot().active).toBe(2); + expect(p.snapshot().maxWaiters).toBe(1); + + a.release(); // hands the slot to c + const c = await cPromise; + expect(cResolved).toBe(true); + expect(p.snapshot().active).toBe(2); // handoff kept the cap, never 3 + + b.release(); + c.release(); + expect(p.snapshot().active).toBe(0); + }); + + test('double release is idempotent', async () => { + const p = pacer({ maxConcurrency: 1 }); + const a = await p.acquire(); + a.release(); + a.release(); + expect(p.snapshot().active).toBe(0); + }); +}); + +describe('abort throws (never falls into a DB call)', () => { + test('acquire throws AbortError when the signal fires while waiting', async () => { + const p = pacer({ maxConcurrency: 1 }); + await p.acquire(); // exhausts the single slot + const ac = new AbortController(); + const blocked = p.acquire(ac.signal); + ac.abort(); + await expect(blocked).rejects.toBeInstanceOf(AbortError); + expect(p.snapshot().maxWaiters).toBe(1); + }); + + test('acquire throws immediately when already aborted', async () => { + const p = pacer(); + const ac = new AbortController(); + ac.abort(); + await expect(p.acquire(ac.signal)).rejects.toBeInstanceOf(AbortError); + }); + + test('pace throws AbortError when aborted during sleep', async () => { + // sleep that rejects with AbortError when the signal is set. + const p = createDbPacer({ + bundle: BUNDLE, + sleep: (_ms, signal) => + signal?.aborted ? Promise.reject(new AbortError('aborted')) : Promise.resolve(), + rng: () => 0.5, + }); + p.observe(1000); // EWMA well above paceAtMs → will sleep + const ac = new AbortController(); + ac.abort(); + await expect(p.pace(ac.signal)).rejects.toBeInstanceOf(AbortError); + }); +}); + +describe('in-band EWMA + observe guard', () => { + test('first sample seeds EWMA; later samples smooth it', () => { + const p = pacer({ ewmaAlpha: 0.5 }); + expect(p.snapshot().ewmaMs).toBeNull(); + p.observe(200); + expect(p.snapshot().ewmaMs).toBe(200); + p.observe(400); + expect(p.snapshot().ewmaMs).toBe(300); // 0.5*400 + 0.5*200 + expect(p.snapshot().sampleCount).toBe(2); + }); + + test('NaN / negative samples are ignored', () => { + const p = pacer(); + p.observe(Number.NaN); + p.observe(-5); + p.observe(Infinity); + expect(p.snapshot().ewmaMs).toBeNull(); + expect(p.snapshot().sampleCount).toBe(0); + }); +}); + +describe('cooperative sleep', () => { + test('no sleep when EWMA is below paceAtMs', async () => { + const p = pacer({ paceAtMs: 100 }); + p.observe(50); + await p.pace(); + expect(p.snapshot().sleepCount).toBe(0); + expect(p.snapshot().totalSleptMs).toBe(0); + }); + + test('sleeps when EWMA exceeds paceAtMs; jittered to 50-100% of base, capped', async () => { + // rng=0 → jitter factor 0.5 → ms = round(base * 0.5). + const p = pacer({ paceAtMs: 100, maxSleepMs: 1000 }, () => 0); + p.observe(800); // base = min(1000, 800) = 800; ms = 400 + await p.pace(); + expect(p.snapshot().sleepCount).toBe(1); + expect(p.snapshot().totalSleptMs).toBe(400); + }); + + test('sleep base is capped at maxSleepMs', async () => { + const p = pacer({ paceAtMs: 100, maxSleepMs: 300 }, () => 1); // jitter factor 1.0 + p.observe(5000); // base = min(300, 5000) = 300; ms = 300 + await p.pace(); + expect(p.snapshot().totalSleptMs).toBe(300); + }); + + test('different rng values decorrelate sleep durations (anti-thundering-herd)', async () => { + const lo = pacer({ paceAtMs: 100 }, () => 0); // factor 0.5 + const hi = pacer({ paceAtMs: 100 }, () => 1); // factor 1.0 + lo.observe(800); + hi.observe(800); + await lo.pace(); + await hi.pace(); + expect(lo.snapshot().totalSleptMs).toBe(400); + expect(hi.snapshot().totalSleptMs).toBe(800); + }); +}); + +describe('fail-open', () => { + test('an unexpected (non-abort) sleep failure does not throw', async () => { + const p = createDbPacer({ + bundle: BUNDLE, + sleep: () => Promise.reject(new Error('boom')), + rng: () => 0.5, + }); + p.observe(1000); + await expect(p.pace()).resolves.toBeUndefined(); + // The failed sleep is not counted. + expect(p.snapshot().sleepCount).toBe(0); + }); +}); + +describe('no-op pacer (off mode / PGLite)', () => { + test('createNoopPacer: unbounded acquire, zero sleeps', async () => { + const p = createNoopPacer(); + const permits = await Promise.all([p.acquire(), p.acquire(), p.acquire()]); + expect(permits).toHaveLength(3); + p.observe(99999); + await p.pace(); + expect(p.snapshot().enabled).toBe(false); + expect(p.snapshot().totalSleptMs).toBe(0); + }); + + test('an off bundle yields a no-op pacer', async () => { + const p = createDbPacer({ bundle: PACE_BUNDLES.off }); + expect(p.snapshot().enabled).toBe(false); + await p.acquire(); + await p.acquire(); // never blocks despite maxConcurrency 0 + expect(p.snapshot().active).toBe(0); + }); +}); + +describe('dispose', () => { + test('dispose releases blocked acquirers with a no-op permit', async () => { + const p = pacer({ maxConcurrency: 1 }); + await p.acquire(); + const blocked = p.acquire(); + p.dispose(); + const permit = await blocked; // resolves instead of hanging + expect(permit).toBeDefined(); + }); +}); diff --git a/test/pace-mode.test.ts b/test/pace-mode.test.ts new file mode 100644 index 000000000..414662d6f --- /dev/null +++ b/test/pace-mode.test.ts @@ -0,0 +1,120 @@ +/** + * Pins the DB-pacing mode bundles + resolution chain. + * The load-bearing claim: env beats config (incident escape hatch), per-call + * beats env, and `off` resolves to a disabled bundle. + */ +import { describe, expect, test } from 'bun:test'; +import { + PACE_BUNDLES, + PACE_MODES, + DEFAULT_PACE_MODE, + isPaceMode, + resolvePaceMode, + loadOverridesFromConfig, + readPaceEnv, +} from '../src/core/pace-mode.ts'; + +describe('pace-mode bundles', () => { + test('off is the default and is disabled', () => { + expect(DEFAULT_PACE_MODE).toBe('off'); + expect(PACE_BUNDLES.off.enabled).toBe(false); + }); + + test('enabled bundles cap concurrency below the unpaced default (20)', () => { + for (const m of ['gentle', 'balanced', 'aggressive'] as const) { + expect(PACE_BUNDLES[m].enabled).toBe(true); + expect(PACE_BUNDLES[m].maxConcurrency).toBeGreaterThanOrEqual(1); + expect(PACE_BUNDLES[m].maxConcurrency).toBeLessThan(20); + } + }); + + test('isPaceMode guards the union', () => { + expect(PACE_MODES.every(isPaceMode)).toBe(true); + expect(isPaceMode('turbo')).toBe(false); + expect(isPaceMode(undefined)).toBe(false); + }); +}); + +describe('resolvePaceMode precedence', () => { + test('unknown mode falls back to off (disabled), mode_valid=false', () => { + const r = resolvePaceMode({ mode: 'turbo' }); + expect(r.resolved_mode).toBe('off'); + expect(r.mode_valid).toBe(false); + expect(r.enabled).toBe(false); + }); + + test('config mode applies its bundle', () => { + const r = resolvePaceMode({ mode: 'balanced' }); + expect(r.resolved_mode).toBe('balanced'); + expect(r.maxConcurrency).toBe(PACE_BUNDLES.balanced.maxConcurrency); + }); + + test('env mode beats config mode (incident escape hatch)', () => { + const r = resolvePaceMode({ mode: 'gentle', envMode: 'aggressive' }); + expect(r.resolved_mode).toBe('aggressive'); + }); + + test('per-call mode beats env and config', () => { + const r = resolvePaceMode({ mode: 'gentle', envMode: 'balanced', perCallMode: 'aggressive' }); + expect(r.resolved_mode).toBe('aggressive'); + }); + + test('env knob override beats config knob override', () => { + const r = resolvePaceMode({ + mode: 'balanced', + configOverrides: { maxConcurrency: 6 }, + envOverrides: { maxConcurrency: 2 }, + }); + expect(r.maxConcurrency).toBe(2); + }); + + test('per-call knob beats env and config', () => { + const r = resolvePaceMode({ + mode: 'balanced', + configOverrides: { maxConcurrency: 6 }, + envOverrides: { maxConcurrency: 2 }, + perCall: { maxConcurrency: 12 }, + }); + expect(r.maxConcurrency).toBe(12); + }); + + test('clamps an out-of-range maxConcurrency back to the bundle when enabled', () => { + const r = resolvePaceMode({ mode: 'balanced', perCall: { maxConcurrency: 0 } }); + expect(r.enabled).toBe(true); + expect(r.maxConcurrency).toBe(PACE_BUNDLES.balanced.maxConcurrency); + }); + + test('clamps a bad ewmaAlpha back to a sane default', () => { + const r = resolvePaceMode({ mode: 'balanced', perCall: { ewmaAlpha: 5 } }); + expect(r.ewmaAlpha).toBeGreaterThan(0); + expect(r.ewmaAlpha).toBeLessThanOrEqual(1); + }); +}); + +describe('config + env parsing', () => { + test('loadOverridesFromConfig parses present keys only', () => { + const ov = loadOverridesFromConfig({ + 'pace.max_concurrency': '5', + 'pace.pace_at_ms': '400', + }); + expect(ov.maxConcurrency).toBe(5); + expect(ov.paceAtMs).toBe(400); + expect(ov.maxSleepMs).toBeUndefined(); + }); + + test('readPaceEnv reads GBRAIN_PACE_* including mode', () => { + const { envMode, envOverrides } = readPaceEnv({ + GBRAIN_PACE_MODE: 'gentle', + GBRAIN_PACE_MAX_CONCURRENCY: '3', + GBRAIN_PACE_ENABLED: 'true', + }); + expect(envMode).toBe('gentle'); + expect(envOverrides.maxConcurrency).toBe(3); + expect(envOverrides.enabled).toBe(true); + }); + + test('rejects out-of-range env concurrency (falls through)', () => { + const { envOverrides } = readPaceEnv({ GBRAIN_PACE_MAX_CONCURRENCY: '99999' }); + expect(envOverrides.maxConcurrency).toBeUndefined(); + }); +});