diff --git a/src/core/postgres-engine.ts b/src/core/postgres-engine.ts index b42f2a909..146aea41d 100644 --- a/src/core/postgres-engine.ts +++ b/src/core/postgres-engine.ts @@ -5565,30 +5565,78 @@ export class PostgresEngine implements BrainEngine { }); } + /** + * perf (#1694 by @Omerbahari): process-lifetime config cache. A single + * search fires ~85 getConfig() reads (loadConfigWithEngine x2, plus + * mode/cache/intent/rerank/graph-signals resolvers). On a remote pooler + * each read is a round-trip; serial they dominate query latency and can + * push the op handler past cli.ts's 10s disconnect force-exit, truncating + * stdout. First read batch-loads the whole `config` table into this Map + * (inside the same connRetry posture as the per-key read — #1603/#1891); + * setConfig/unsetConfig write through. TTL bounds staleness for + * multi-writer processes; GBRAIN_CONFIG_CACHE_TTL_MS=0 disables. + * Only present keys are stored — Map.has() distinguishes known-absent. + */ + private _configCache: Map | null = null; + private _configCacheLoadedAt = 0; + private _configCacheLoad: Promise | null = null; + private get _configCacheTtlMs(): number { + const raw = process.env.GBRAIN_CONFIG_CACHE_TTL_MS; + if (raw !== undefined) { + const n = parseInt(raw, 10); + if (Number.isFinite(n) && n >= 0) return n; + } + return 30_000; + } + async getConfig(key: string): Promise { // #1603: a transient pooler drop on this read used to throw / fall through // to defaults silently — which on remote Postgres surfaces as the wrong - // search mode/knobs and empty-stdout queries. - return this.connRetry(async () => { - const rows = await this.sql`SELECT value FROM config WHERE key = ${key}`; - return rows.length > 0 ? (rows[0].value as string) : null; - }); + // search mode/knobs and empty-stdout queries. Both the batch load and the + // cache-off per-key read keep the connRetry reconnect posture. + const ttl = this._configCacheTtlMs; + if (ttl === 0) { + return this.connRetry(async () => { + const rows = await this.sql`SELECT value FROM config WHERE key = ${key}`; + return rows.length > 0 ? (rows[0].value as string) : null; + }); + } + if (this._configCache === null || Date.now() - this._configCacheLoadedAt >= ttl) { + // Single-flight: concurrent cold reads share one batch load. + this._configCacheLoad ??= this.connRetry(async () => { + const rows = await this.sql`SELECT key, value FROM config` as unknown as + Array<{ key: string; value: string | null }>; + const map = new Map(); + for (const r of rows) if (r.value != null) map.set(r.key, r.value); + this._configCache = map; + this._configCacheLoadedAt = Date.now(); + }).finally(() => { + this._configCacheLoad = null; + }); + await this._configCacheLoad; + } + return this._configCache!.has(key) ? this._configCache!.get(key)! : null; } async setConfig(key: string, value: string): Promise { - return this.connRetry(async () => { + await this.connRetry(async () => { await this.sql` INSERT INTO config (key, value) VALUES (${key}, ${value}) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value `; }); + // Write-through so a long-lived process never serves stale config. + this._configCache?.set(key, value); } async unsetConfig(key: string): Promise { - return this.connRetry(async () => { + const count = await this.connRetry(async () => { const result = await this.sql`DELETE FROM config WHERE key = ${key}` as unknown as { count: number }; return result.count ?? 0; }); + // Write-through: known-absent, so the cache doesn't serve a stale value. + this._configCache?.delete(key); + return count; } async listConfigKeys(prefix: string): Promise { diff --git a/test/e2e/helpers.ts b/test/e2e/helpers.ts index 0d0fe762b..84685451e 100644 --- a/test/e2e/helpers.ts +++ b/test/e2e/helpers.ts @@ -29,6 +29,13 @@ if (existsSync(envPath)) { } } +// E2E suites seed/rewrite the config table via raw SQL and expect engine +// reads to see it immediately; disable the process-lifetime config cache +// (#1694) so read semantics match pre-cache behavior. Spawned CLI +// subprocesses inherit this. Cache semantics are pinned by +// test/postgres-engine-config-cache.test.ts. +process.env.GBRAIN_CONFIG_CACHE_TTL_MS ??= '0'; + const DATABASE_URL = process.env.DATABASE_URL; const FIXTURES_DIR = resolve(import.meta.dir, 'fixtures'); diff --git a/test/postgres-engine-config-cache.test.ts b/test/postgres-engine-config-cache.test.ts new file mode 100644 index 000000000..949034a10 --- /dev/null +++ b/test/postgres-engine-config-cache.test.ts @@ -0,0 +1,120 @@ +/** + * Process-lifetime config cache (#1694 takeover, by @Omerbahari). + * + * A single search fires ~85 getConfig() reads; on a remote pooler each is a + * round-trip. The first read now batch-loads the whole `config` table into a + * Map; setConfig/unsetConfig write through; TTL bounds multi-writer + * staleness; GBRAIN_CONFIG_CACHE_TTL_MS=0 restores per-key reads. + * + * Pure: stubs `_sql` with a call-counting fake; no real DB. + */ + +import { describe, it, expect, beforeEach, afterEach } from 'bun:test'; +import { PostgresEngine } from '../src/core/postgres-engine.ts'; + +const FAST_RETRY = { maxRetries: 3, delayMs: 1, delayMaxMs: 1, jitter: 'none' as const }; + +/** Engine whose `sql` records every query's template strings and returns `rows`. */ +function makeEngine(rows: unknown[]) { + const e = new PostgresEngine(); + const calls: string[] = []; + (e as unknown as { _connectionStyle: string })._connectionStyle = 'instance'; + (e as unknown as { _bulkRetryOptsCache: unknown })._bulkRetryOptsCache = FAST_RETRY; + (e as unknown as { _sql: unknown })._sql = (strings: TemplateStringsArray) => { + calls.push(strings.join('?')); + return Promise.resolve(rows); + }; + return { engine: e, calls }; +} + +const savedTtl = process.env.GBRAIN_CONFIG_CACHE_TTL_MS; + +beforeEach(() => { + delete process.env.GBRAIN_CONFIG_CACHE_TTL_MS; +}); + +afterEach(() => { + if (savedTtl === undefined) delete process.env.GBRAIN_CONFIG_CACHE_TTL_MS; + else process.env.GBRAIN_CONFIG_CACHE_TTL_MS = savedTtl; +}); + +describe('PostgresEngine config cache (#1694)', () => { + it('batch-loads once and serves repeat reads from the cache', async () => { + const { engine, calls } = makeEngine([ + { key: 'search.mode', value: 'balanced' }, + { key: 'embedding_multimodal', value: 'true' }, + ]); + expect(await engine.getConfig('search.mode')).toBe('balanced'); + expect(await engine.getConfig('embedding_multimodal')).toBe('true'); + expect(await engine.getConfig('search.mode')).toBe('balanced'); + // One SELECT total — this is the whole point of the fix. + expect(calls.length).toBe(1); + expect(calls[0]).toContain('SELECT key, value FROM config'); + }); + + it('returns null for a known-absent key without an extra round-trip', async () => { + const { engine, calls } = makeEngine([{ key: 'a', value: '1' }]); + expect(await engine.getConfig('missing.key')).toBeNull(); + expect(await engine.getConfig('missing.key')).toBeNull(); + expect(calls.length).toBe(1); + }); + + it('setConfig writes through so subsequent reads see the new value', async () => { + const { engine, calls } = makeEngine([{ key: 'k', value: 'old' }]); + expect(await engine.getConfig('k')).toBe('old'); + await engine.setConfig('k', 'new'); + expect(await engine.getConfig('k')).toBe('new'); + expect(calls.length).toBe(2); // batch load + upsert; no re-read + }); + + it('unsetConfig writes through so subsequent reads see absence', async () => { + const { engine } = makeEngine([{ key: 'k', value: 'v' }]); + expect(await engine.getConfig('k')).toBe('v'); + await engine.unsetConfig('k'); + expect(await engine.getConfig('k')).toBeNull(); + }); + + it('concurrent cold reads share a single batch load (single-flight)', async () => { + const { engine, calls } = makeEngine([{ key: 'k', value: 'v' }]); + const [a, b, c] = await Promise.all([ + engine.getConfig('k'), + engine.getConfig('k'), + engine.getConfig('other'), + ]); + expect([a, b, c]).toEqual(['v', 'v', null]); + expect(calls.length).toBe(1); + }); + + it('GBRAIN_CONFIG_CACHE_TTL_MS=0 disables the cache (per-key reads)', async () => { + process.env.GBRAIN_CONFIG_CACHE_TTL_MS = '0'; + const { engine, calls } = makeEngine([{ value: 'v' }]); + expect(await engine.getConfig('k')).toBe('v'); + expect(await engine.getConfig('k')).toBe('v'); + expect(calls.length).toBe(2); + expect(calls[0]).toContain('SELECT value FROM config WHERE key ='); + }); + + it('an expired TTL reloads from the database', async () => { + process.env.GBRAIN_CONFIG_CACHE_TTL_MS = '1'; + const { engine, calls } = makeEngine([{ key: 'k', value: 'v' }]); + expect(await engine.getConfig('k')).toBe('v'); + await new Promise((r) => setTimeout(r, 5)); + expect(await engine.getConfig('k')).toBe('v'); + expect(calls.length).toBe(2); // two batch loads + }); + + it('the batch load keeps the connRetry reconnect posture (#1603/#1891)', async () => { + const e = new PostgresEngine(); + (e as unknown as { _connectionStyle: string })._connectionStyle = 'instance'; + (e as unknown as { _sql: unknown })._sql = null; // torn-down pool → retryable + (e as unknown as { _bulkRetryOptsCache: unknown })._bulkRetryOptsCache = FAST_RETRY; + let reconnects = 0; + (e as unknown as { reconnect: () => Promise }).reconnect = async () => { + reconnects++; + (e as unknown as { _sql: unknown })._sql = () => + Promise.resolve([{ key: 'k', value: 'v' }]); + }; + expect(await e.getConfig('k')).toBe('v'); + expect(reconnects).toBe(1); + }); +}); diff --git a/test/postgres-engine-config-reconnect.test.ts b/test/postgres-engine-config-reconnect.test.ts index df44d47c9..a4770f845 100644 --- a/test/postgres-engine-config-reconnect.test.ts +++ b/test/postgres-engine-config-reconnect.test.ts @@ -43,7 +43,9 @@ function makeTornDownEngine(poolResult: unknown): { engine: PostgresEngine; reco describe('PostgresEngine non-batch config accessors self-heal (PR #1891 takeover)', () => { it('getConfig reconnects + retries a null instance pool, then returns the value', async () => { - const { engine, reconnects } = makeTornDownEngine([{ value: 'live-value' }]); + // Rows carry `key` too: getConfig's default cached path batch-loads + // `SELECT key, value FROM config` (#1694) through the same connRetry. + const { engine, reconnects } = makeTornDownEngine([{ key: 'some.key', value: 'live-value' }]); expect(await engine.getConfig('some.key')).toBe('live-value'); expect(reconnects()).toBe(1); // exactly one reconnect closed the gap });