mirror of
https://github.com/garrytan/gbrain.git
synced 2026-07-29 19:01:39 +00:00
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.
This commit is contained in:
Garry Tan
@@ -0,0 +1,304 @@
|
||||
/**
|
||||
* db-pacer.ts — composable DB-contention pacing primitive (paced-backfill internals).
|
||||
*
|
||||
* gbrain's bulk write paths (embed --stale, sync) can saturate a PgBouncer
|
||||
* transaction-mode pooler and starve the minion supervisor's lock renewals.
|
||||
* An operator's external SIGSTOP/SIGCONT wrapper proved the idea but was blind
|
||||
* (out-of-band probe pool), unsafe (froze the child mid-transaction), and
|
||||
* couldn't touch PEAK pressure (it paused between batches while N workers were
|
||||
* already in flight).
|
||||
*
|
||||
* This primitive fixes all three:
|
||||
* - CONCURRENCY CAP is the real lever. `acquire()` is a counting semaphore
|
||||
* that bounds simultaneous in-flight DB writes — directly limiting pooler
|
||||
* slots held at once. (Single-pool callers like embed instead lower their
|
||||
* worker count to `maxConcurrency`; the permit exists for the MULTI-pool
|
||||
* case — sync's separate engine per worker — where one budget must span
|
||||
* pools a single worker-count can't.)
|
||||
* - IN-BAND latency is the signal. `observe(ms)` feeds an EWMA from the work's
|
||||
* OWN queries — it can never be blind the way a separate probe pool was.
|
||||
* - COOPERATIVE sleep, never SIGSTOP. `pace()` sleeps between safe points on
|
||||
* `setTimeout` (timers phase), so the lock-heartbeat `setInterval` keeps
|
||||
* firing. Per-call jitter prevents a 20-worker thundering-herd resume.
|
||||
*
|
||||
* Contracts:
|
||||
* - `acquire()` / `pace()` THROW `AbortError` if the signal fires while
|
||||
* waiting — a cancel can never fall through into a DB call.
|
||||
* - FAIL-OPEN: any unexpected internal error degrades to a no-op; a pacer bug
|
||||
* must never kill a backfill, and nothing here throws an unhandledRejection
|
||||
* (the failure class behind the prior #1972 lock-renewal crash).
|
||||
* - `off` mode / PGLite → `createNoopPacer()`: unbounded permits, zero sleeps.
|
||||
*/
|
||||
|
||||
import { AbortError, anySignal } from './abort-check.ts';
|
||||
import type { PaceBundle } from './pace-mode.ts';
|
||||
|
||||
/** Snapshot of pacer state for telemetry + tests. */
|
||||
export interface PaceSnapshot {
|
||||
enabled: boolean;
|
||||
maxConcurrency: number;
|
||||
/** Permits currently held. */
|
||||
active: number;
|
||||
/** EWMA of observed DB-op latency (ms); null until the first sample. */
|
||||
ewmaMs: number | null;
|
||||
/** Cumulative ms spent in cooperative sleeps. */
|
||||
totalSleptMs: number;
|
||||
/** Number of `pace()` calls that actually slept. */
|
||||
sleepCount: number;
|
||||
/** High-water mark of waiters blocked in `acquire()`. */
|
||||
maxWaiters: number;
|
||||
/** Count of `observe()` samples folded into the EWMA. */
|
||||
sampleCount: number;
|
||||
}
|
||||
|
||||
/** A held permit. `release()` is idempotent. */
|
||||
export interface Permit {
|
||||
release(): void;
|
||||
}
|
||||
|
||||
export interface DbPacer {
|
||||
/**
|
||||
* Acquire a DB-write permit (caps simultaneous in-flight writes to
|
||||
* `maxConcurrency`). THROWS `AbortError` if `signal` fires while waiting.
|
||||
* On a disabled pacer, resolves immediately with a no-op permit.
|
||||
*/
|
||||
acquire(signal?: AbortSignal): Promise<Permit>;
|
||||
/** 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<void>;
|
||||
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<void>;
|
||||
/** 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<void> {
|
||||
if (signal?.aborted) return Promise.reject(new AbortError(abortReason(signal)));
|
||||
if (!(ms > 0)) return Promise.resolve();
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
const timer = setTimeout(() => {
|
||||
cleanup();
|
||||
resolve();
|
||||
}, ms);
|
||||
// Don't keep the event loop alive purely for a pacing sleep.
|
||||
(timer as { unref?: () => void }).unref?.();
|
||||
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<Permit> {
|
||||
try {
|
||||
if (signal?.aborted) throw new AbortError(abortReason(signal));
|
||||
if (disposed) return NOOP_PERMIT;
|
||||
if (active < max) {
|
||||
active++;
|
||||
return makePermit();
|
||||
}
|
||||
return await new Promise<Permit>((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<void> {
|
||||
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<T>(pacer: DbPacer, fn: () => Promise<T>): Promise<T> {
|
||||
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 };
|
||||
@@ -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;
|
||||
@@ -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<PaceMode> = 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<Record<PaceMode, Readonly<PaceBundle>>> = 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<string>).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=<mode>` (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 = <K extends keyof PaceBundle>(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<string> = 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<string, string | undefined>,
|
||||
): 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<string, string | undefined> = 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<keyof PaceKeyOverrides, string>,
|
||||
): 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<string | null>;
|
||||
}): Promise<{ mode?: string; configOverrides: PaceKeyOverrides }> {
|
||||
const safeGet = async (k: string): Promise<string | undefined> => {
|
||||
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<string, string | undefined> = {};
|
||||
PACE_MODE_CONFIG_KEYS.forEach((key, i) => {
|
||||
if (vals[i] !== undefined) configMap[key] = vals[i];
|
||||
});
|
||||
return { mode, configOverrides: loadOverridesFromConfig(configMap) };
|
||||
}
|
||||
@@ -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<void> {
|
||||
if (signal?.aborted) return Promise.reject(new AbortError('aborted'));
|
||||
return Promise.resolve();
|
||||
}
|
||||
|
||||
function pacer(over: Partial<PaceBundle> = {}, 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();
|
||||
});
|
||||
});
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user