fix(db-lock): heartbeat-aware takeover + direct-pool refresh (#1794)

src/core/db-lock.ts

@@ -42,6 +42,30 @@ const DEFAULT_TTL_MINUTES = 30;
  */
 export const HOLDER_TAKEOVER_GRACE_MS = 60_000;
 
+/**
+ * v0.42.x (#1794): heartbeat-aware steal grace. A holder whose
+ * `last_refreshed_at` is within this window is treated as ALIVE and is NOT
+ * stolen even if its `ttl_expires_at` has lapsed — defending a live, actively
+ * refreshing holder whose refresh tick was briefly starved (the #1794 thrash,
+ * where a CPU-bound import let the TTL expire and a competing launch stole the
+ * live lock). A genuinely dead holder stops refreshing, ages past the grace,
+ * and becomes stealable again (TTL stays the ultimate backstop). Derived from
+ * the TTL so it scales with the refresh cadence; override with
+ * GBRAIN_LOCK_STEAL_GRACE_SECONDS.
+ */
+export const DEFAULT_STEAL_GRACE_SECONDS = 600;
+
+export function resolveStealGraceSeconds(ttlMinutes: number): number {
+  const raw = process.env.GBRAIN_LOCK_STEAL_GRACE_SECONDS;
+  if (raw) {
+    const n = Number(raw);
+    if (Number.isInteger(n) && n > 0) return n;
+  }
+  // Refresh fires ~ttl/6; protect a holder that refreshed within ~2 ticks.
+  const refreshSec = Math.max(15, (ttlMinutes * 60) / 6);
+  return Math.max(Math.floor(refreshSec * 2), 60);
+}
+
 /**
  * Liveness classification of a lock holder, from the perspective of the
  * current host. Shared by `isHolderDeadLocally` (auto-takeover in
@@ -130,6 +154,9 @@ export async function tryAcquireDbLock(
 ): Promise<DbLockHandle | null> {
   const pid = process.pid;
   const host = hostname();
+  // v0.42.x (#1794): a holder that refreshed within this window is protected
+  // from the ON CONFLICT steal even if its TTL lapsed (starved-but-alive).
+  const stealGraceSeconds = resolveStealGraceSeconds(ttlMinutes);
 
   // Engine-agnostic: prefer the engine's raw escape hatch (`sql` for postgres-js,
   // `db.query` for PGLite). Mirrors cycle.ts's pattern so behavior stays identical.
@@ -166,6 +193,8 @@ export async function tryAcquireDbLock(
             ttl_expires_at = NOW() + ${ttl}::interval,
             last_refreshed_at = NOW()
         WHERE gbrain_cycle_locks.ttl_expires_at < NOW()
+          AND (gbrain_cycle_locks.last_refreshed_at IS NULL
+               OR gbrain_cycle_locks.last_refreshed_at < NOW() - ${stealGraceSeconds} * INTERVAL '1 second')
       RETURNING id
     `;
     if (rows.length === 0) return null;
@@ -179,14 +208,16 @@ export async function tryAcquireDbLock(
       id: lockId,
       refresh: async () => {
         // v0.41.13.0: bump BOTH ttl_expires_at AND last_refreshed_at.
-        // Without last_refreshed_at, --max-age would steal healthy locks
-        // whose acquired_at is old but whose holder is alive and refreshing.
-        await sql`
-          UPDATE gbrain_cycle_locks
-            SET ttl_expires_at = NOW() + ${ttl}::interval,
-                last_refreshed_at = NOW()
-          WHERE id = ${lockId} AND holder_pid = ${pid}
-        `;
+        // v0.42.x (#1794): route through the DIRECT session pool, not the
+        // transaction pool, so a Supavisor pooler exhaustion (EMAXCONNSESSION)
+        // can't kill the heartbeat and let the live lock get stolen.
+        await engine.executeRawDirect(
+          `UPDATE gbrain_cycle_locks
+              SET ttl_expires_at = NOW() + ($1)::interval,
+                  last_refreshed_at = NOW()
+            WHERE id = $2 AND holder_pid = $3`,
+          [ttl, lockId, pid],
+        );
       },
       release: async () => {
         deregister();
@@ -211,8 +242,10 @@ export async function tryAcquireDbLock(
              ttl_expires_at = NOW() + $4::interval,
              last_refreshed_at = NOW()
          WHERE gbrain_cycle_locks.ttl_expires_at < NOW()
+           AND (gbrain_cycle_locks.last_refreshed_at IS NULL
+                OR gbrain_cycle_locks.last_refreshed_at < NOW() - $5 * INTERVAL '1 second')
        RETURNING id`,
-      [lockId, pid, host, ttl],
+      [lockId, pid, host, ttl, stealGraceSeconds],
     );
     if (rows.length === 0) return null;
     const deregister = registerCleanup(`db-lock:${lockId}`, async () => {
@@ -648,27 +681,25 @@ export async function withRefreshingLock<T>(
   const interval = setInterval(() => {
     void (async () => {
       try {
-        // A4 heartbeat: SELECT 1 against the engine's connection pool.
-        // Honest limit: this checks a connection is responsive in general,
-        // not the SPECIFIC backend running `work()`. The full X1 fix
-        // (lock-refresh on the work-pinned connection via withReservedConnection)
-        // is layered in by callers that pass the work backend's sql in.
-        // For migrate.ts (transactional DDL), the engine.transaction() path
-        // pins the backend; the heartbeat against engine.sql is a useful
-        // proxy for "Postgres is reachable" even if it can race the actual
-        // backend's wedge state. Lane B's primary win is the auto-refresh
-        // itself; the precise-backend-bind heartbeat is a Lane B follow-up.
-        const probe = engineSelectOne(engine);
+        // v0.42.x (#1794, V1): the refresh IS the heartbeat. handle.refresh()
+        // routes through the DIRECT session pool (postgres), so it survives a
+        // transaction-pool exhaustion (EMAXCONNSESSION) that would otherwise
+        // kill renewal and let the live lock be stolen. The pre-v0.42 code first
+        // probed `SELECT 1` on the READ pool and clearInterval'd on probe
+        // failure — that's exactly how an exhausted read pool stopped renewal
+        // even though the lock was alive. We no longer gate renewal on read-pool
+        // health, and we do NOT clearInterval on a transient failure: a blip
+        // self-heals on the next tick; the TTL is the backstop if the pool stays
+        // genuinely dead (at which point a steal is correct).
         const timeout = new Promise<never>((_, reject) =>
-          setTimeout(() => reject(new Error('heartbeat_timeout')), heartbeatTimeoutMs)
+          setTimeout(() => reject(new Error('refresh_timeout')), heartbeatTimeoutMs)
         );
-        await Promise.race([probe, timeout]);
-        await handle.refresh();
+        await Promise.race([handle.refresh(), timeout]);
+        healthOk = true;
       } catch (err) {
         const msg = err instanceof Error ? err.message : String(err);
-        process.stderr.write(`[lock-refresh] ${lockId}: ${msg}; lock will auto-expire\n`);
+        process.stderr.write(`[lock-refresh] ${lockId}: ${msg}; will retry next tick\n`);
         healthOk = false;
-        clearInterval(interval);
       }
     })();
   }, refreshIntervalMs);
@@ -690,24 +721,6 @@ export async function withRefreshingLock<T>(
   }
 }
 
-/** Internal: SELECT 1 on the engine's connection. */
-async function engineSelectOne(engine: BrainEngine): Promise<void> {
-  const maybePG = engine as unknown as { sql?: (...args: unknown[]) => Promise<unknown> };
-  const maybePGLite = engine as unknown as {
-    db?: { query: (sql: string) => Promise<{ rows: unknown[] }> };
-  };
-  if (engine.kind === 'postgres' && maybePG.sql) {
-    const sql = maybePG.sql as any;
-    await sql`SELECT 1`;
-    return;
-  }
-  if (engine.kind === 'pglite' && maybePGLite.db) {
-    await maybePGLite.db.query('SELECT 1');
-    return;
-  }
-  throw new Error(`Unknown engine kind for heartbeat: ${engine.kind}`);
-}
-
 /**
  * v0.41 Eng D9 (codex pass-2 #7 + #8) — per-tick election convenience.
  *

test/db-lock-heartbeat-takeover.test.ts

@@ -0,0 +1,162 @@
+/**
+ * #1794 — heartbeat-aware lock takeover + direct-pool refresh.
+ *
+ * The lock-thrash fix: a holder that refreshed within the steal-grace window is
+ * NOT stolen even if its TTL lapsed (starved-but-alive), while a holder that
+ * stopped refreshing past the grace IS stolen. These tests isolate the ON
+ * CONFLICT grace logic by holding with `process.pid` (alive) so the dead-pid
+ * auto-takeover path can't fire — a successful steal therefore proves the
+ * grace/last_refreshed_at predicate did it.
+ *
+ * Plus the commit-8 causal-mechanism check: setImmediate yields let a
+ * setInterval tick fire during a busy loop (the reason the import loop's
+ * maybeYield keeps the refresh heartbeat alive).
+ */
+
+import { describe, test, expect, beforeAll, afterAll, beforeEach } from 'bun:test';
+import { hostname } from 'os';
+import { PGLiteEngine } from '../src/core/pglite-engine.ts';
+import {
+  tryAcquireDbLock,
+  resolveStealGraceSeconds,
+  DEFAULT_STEAL_GRACE_SECONDS,
+} from '../src/core/db-lock.ts';
+
+let engine: PGLiteEngine;
+
+beforeAll(async () => {
+  engine = new PGLiteEngine();
+  await engine.connect({ database_url: '' });
+  await engine.initSchema();
+});
+
+afterAll(async () => {
+  await engine.disconnect();
+});
+
+beforeEach(async () => {
+  await engine.executeRaw(`DELETE FROM gbrain_cycle_locks WHERE id LIKE 'test-hb-%'`);
+});
+
+const LOCAL = hostname();
+
+/**
+ * Seed a TTL-EXPIRED lock held by an ALIVE pid (process.pid), with a chosen
+ * last_refreshed_at age (or NULL). Alive holder → dead-pid auto-takeover can't
+ * fire, so the only reclaim path is the ON CONFLICT grace predicate.
+ */
+async function seedExpiredLock(id: string, refreshedSecondsAgo: number | null): Promise<void> {
+  if (refreshedSecondsAgo === null) {
+    await engine.executeRaw(
+      `INSERT INTO gbrain_cycle_locks (id, holder_pid, holder_host, acquired_at, ttl_expires_at, last_refreshed_at)
+       VALUES ($1, $2, $3, NOW() - INTERVAL '1 hour', NOW() - INTERVAL '5 minutes', NULL)`,
+      [id, process.pid, LOCAL],
+    );
+  } else {
+    await engine.executeRaw(
+      `INSERT INTO gbrain_cycle_locks (id, holder_pid, holder_host, acquired_at, ttl_expires_at, last_refreshed_at)
+       VALUES ($1, $2, $3, NOW() - INTERVAL '1 hour', NOW() - INTERVAL '5 minutes', NOW() - ($4 || ' seconds')::interval)`,
+      [id, process.pid, LOCAL, String(refreshedSecondsAgo)],
+    );
+  }
+}
+
+async function refreshedAge(id: string): Promise<Date | null> {
+  const rows = await engine.executeRaw<{ last_refreshed_at: string | null }>(
+    `SELECT last_refreshed_at FROM gbrain_cycle_locks WHERE id = $1`,
+    [id],
+  );
+  const v = rows[0]?.last_refreshed_at ?? null;
+  return v ? new Date(v) : null;
+}
+
+describe('resolveStealGraceSeconds', () => {
+  test('derives ~2 refresh ticks from the TTL (30min → 600s)', () => {
+    // refresh ~ttl/6 = 5min = 300s; *2 = 600s.
+    expect(resolveStealGraceSeconds(30)).toBe(DEFAULT_STEAL_GRACE_SECONDS);
+    expect(resolveStealGraceSeconds(30)).toBe(600);
+  });
+
+  test('floors at 60s for tiny TTLs', () => {
+    expect(resolveStealGraceSeconds(1)).toBe(60);
+  });
+
+  test('env override wins', () => {
+    process.env.GBRAIN_LOCK_STEAL_GRACE_SECONDS = '123';
+    try {
+      expect(resolveStealGraceSeconds(30)).toBe(123);
+    } finally {
+      delete process.env.GBRAIN_LOCK_STEAL_GRACE_SECONDS;
+    }
+  });
+
+  test('bad env override falls back to derived', () => {
+    process.env.GBRAIN_LOCK_STEAL_GRACE_SECONDS = 'nope';
+    try {
+      expect(resolveStealGraceSeconds(30)).toBe(600);
+    } finally {
+      delete process.env.GBRAIN_LOCK_STEAL_GRACE_SECONDS;
+    }
+  });
+});
+
+describe('heartbeat-aware takeover (ON CONFLICT grace predicate)', () => {
+  test('FRESH holder is NOT stolen even with an expired TTL', async () => {
+    // ttl expired 5min ago, but refreshed 1s ago → inside the 600s grace.
+    await seedExpiredLock('test-hb-fresh', 1);
+    const handle = await tryAcquireDbLock(engine, 'test-hb-fresh', 30);
+    expect(handle).toBeNull();
+  });
+
+  test('STALE-refresh holder IS stolen (refresh older than the grace)', async () => {
+    // ttl expired AND last refresh 20min ago (> 600s grace) → stealable, even
+    // though the holder pid is alive (proves the grace path, not auto-takeover).
+    await seedExpiredLock('test-hb-stale', 1200);
+    const handle = await tryAcquireDbLock(engine, 'test-hb-stale', 30);
+    expect(handle).not.toBeNull();
+    await handle!.release();
+  });
+
+  test('NULL last_refreshed_at (pre-v98 row) IS stolen on TTL expiry', async () => {
+    await seedExpiredLock('test-hb-null', null);
+    const handle = await tryAcquireDbLock(engine, 'test-hb-null', 30);
+    expect(handle).not.toBeNull();
+    await handle!.release();
+  });
+
+  test('a reclaimed handle refresh() bumps last_refreshed_at (direct pool path)', async () => {
+    await seedExpiredLock('test-hb-bump', 1200);
+    const handle = await tryAcquireDbLock(engine, 'test-hb-bump', 30);
+    expect(handle).not.toBeNull();
+    const before = await refreshedAge('test-hb-bump');
+    // Small real delay so NOW() advances measurably between acquire and refresh.
+    await new Promise((r) => setTimeout(r, 20));
+    await handle!.refresh();
+    const after = await refreshedAge('test-hb-bump');
+    expect(before).not.toBeNull();
+    expect(after).not.toBeNull();
+    expect(after!.getTime()).toBeGreaterThan(before!.getTime());
+    await handle!.release();
+  });
+});
+
+describe('event-loop yield keeps timers alive (commit 8 mechanism)', () => {
+  test('setTimeout(0) yields let a setInterval heartbeat fire during a busy loop', async () => {
+    let ticks = 0;
+    const iv = setInterval(() => { ticks++; }, 2);
+    try {
+      // Mirror the import loop's maybeYield: setTimeout(0) enters the timers
+      // phase, so the setInterval heartbeat can fire mid-loop. (A setImmediate
+      // loop starves the timers phase in Bun — the reason maybeYield uses
+      // setTimeout, not setImmediate.) Bound by wall-clock so the 2ms interval
+      // has real time to fire.
+      const start = Date.now();
+      while (Date.now() - start < 40) {
+        await new Promise<void>((r) => setTimeout(r, 0));
+      }
+    } finally {
+      clearInterval(iv);
+    }
+    expect(ticks).toBeGreaterThan(0);
+  });
+});