gbrain/test/pglite-lock.test.ts

import { describe, test, expect, beforeEach, afterEach } from 'bun:test';
import { mkdirSync, mkdtempSync, rmSync, existsSync, readFileSync, writeFileSync } from 'fs';
import { join } from 'path';
import { tmpdir } from 'os';
import { acquireLock, releaseLock, type LockHandle } from '../src/core/pglite-lock';

const TEST_DIR = join(tmpdir(), 'gbrain-lock-test-' + process.pid);

describe('pglite-lock', () => {
  beforeEach(() => {
    // Clean up test directory
    if (existsSync(TEST_DIR)) rmSync(TEST_DIR, { recursive: true, force: true });
    mkdirSync(TEST_DIR, { recursive: true });
  });

  afterEach(() => {
    if (existsSync(TEST_DIR)) rmSync(TEST_DIR, { recursive: true, force: true });
  });

  test('acquires and releases lock', async () => {
    const lock = await acquireLock(TEST_DIR);
    expect(lock.acquired).toBe(true);
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(true);

    await releaseLock(lock);
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(false);
  });

  test('creates missing data directory before acquiring lock', async () => {
    const missingDataDir = join(TEST_DIR, 'missing-data-dir');

    const lock = await acquireLock(missingDataDir);
    expect(lock.acquired).toBe(true);
    expect(existsSync(missingDataDir)).toBe(true);
    expect(existsSync(join(missingDataDir, '.gbrain-lock'))).toBe(true);

    await releaseLock(lock);
    expect(existsSync(join(missingDataDir, '.gbrain-lock'))).toBe(false);
  });

  test('prevents concurrent lock acquisition', async () => {
    const lock1 = await acquireLock(TEST_DIR, { timeoutMs: 2000 });
    expect(lock1.acquired).toBe(true);

    // Second lock attempt should timeout
    await expect(acquireLock(TEST_DIR, { timeoutMs: 1000 })).rejects.toThrow(/Timed out/);

    await releaseLock(lock1);
  });

  test('detects and cleans stale lock from dead process', async () => {
    // Simulate a stale lock from a dead process
    const lockDir = join(TEST_DIR, '.gbrain-lock');
    mkdirSync(lockDir);
    writeFileSync(join(lockDir, 'lock'), JSON.stringify({
      pid: 999999999, // Non-existent PID
      acquired_at: Date.now(),
      command: 'test',
    }));

    // Should clean up the stale lock and acquire
    const lock = await acquireLock(TEST_DIR);
    expect(lock.acquired).toBe(true);

    await releaseLock(lock);
  });

  test('skips lock for in-memory (undefined dataDir)', async () => {
    const lock = await acquireLock(undefined);
    expect(lock.acquired).toBe(true);
    expect(lock.lockDir).toBe('');

    // Release should be a no-op
    await releaseLock(lock);
  });

  test('lock file contains PID and command', async () => {
    const lock = await acquireLock(TEST_DIR);
    const lockData = JSON.parse(readFileSync(join(TEST_DIR, '.gbrain-lock', 'lock'), 'utf-8'));

    expect(lockData.pid).toBe(process.pid);
    expect(lockData.acquired_at).toBeDefined();
    expect(lockData.command).toBeDefined();

    await releaseLock(lock);
  });

  test('releases lock on disconnect even if DB close fails', async () => {
    const lock = await acquireLock(TEST_DIR);
    expect(lock.acquired).toBe(true);

    // Simulate DB already closed
    await releaseLock(lock);
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(false);

    // Second acquisition should work
    const lock2 = await acquireLock(TEST_DIR);
    expect(lock2.acquired).toBe(true);
    await releaseLock(lock2);
  });
});

describe('pglite-lock #2058 heartbeat + steal-grace', () => {
  beforeEach(() => {
    if (existsSync(TEST_DIR)) rmSync(TEST_DIR, { recursive: true, force: true });
    mkdirSync(TEST_DIR, { recursive: true });
  });
  afterEach(() => {
    if (existsSync(TEST_DIR)) rmSync(TEST_DIR, { recursive: true, force: true });
  });

  function writeHolder(fields: {
    pid: number;
    acquiredAgoMs: number;
    refreshedAgoMs: number;
    command?: string;
    subcommand?: string;
  }) {
    const lockDir = join(TEST_DIR, '.gbrain-lock');
    mkdirSync(lockDir, { recursive: true });
    const now = Date.now();
    writeFileSync(join(lockDir, 'lock'), JSON.stringify({
      pid: fields.pid,
      acquired_at: now - fields.acquiredAgoMs,
      refreshed_at: now - fields.refreshedAgoMs,
      command: fields.command ?? 'test holder',
      ...(fields.subcommand === undefined ? {} : { subcommand: fields.subcommand }),
    }));
  }

  test('a live gbrain serve owner with global flags fails fast with a clear explanation', async () => {
    writeHolder({
      pid: process.pid,
      acquiredAgoMs: 60_000,
      refreshedAgoMs: 0,
      command: '/path with spaces/gbrain/src/cli.ts --quiet serve',
      subcommand: 'serve',
    });

    const startedAt = Date.now();
    await expect(acquireLock(TEST_DIR, { timeoutMs: 5_000 })).rejects.toThrow(
      /already open through `gbrain serve`.*Stop `gbrain serve`, then retry this CLI command.*use its MCP tools instead.*will not remove/s,
    );

    expect(Date.now() - startedAt).toBeLessThan(1_000);
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(true);
  });

  test('legacy serve lock metadata is still recognized', async () => {
    writeHolder({
      pid: process.pid,
      acquiredAgoMs: 60_000,
      refreshedAgoMs: 0,
      command: '/path/to/gbrain/src/cli.ts serve',
    });

    await expect(acquireLock(TEST_DIR, { timeoutMs: 5_000 })).rejects.toThrow(
      /already open through `gbrain serve`/,
    );
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(true);
  });

  test('a search for the word serve is not mistaken for the MCP server', async () => {
    writeHolder({
      pid: process.pid,
      acquiredAgoMs: 60_000,
      refreshedAgoMs: 0,
      command: '/compiled/gbrain search serve',
      subcommand: 'search',
    });

    await expect(acquireLock(TEST_DIR, { timeoutMs: 100 })).rejects.toThrow(/Timed out/);
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(true);
  });

  test('a dead gbrain serve owner is still cleaned up automatically', async () => {
    writeHolder({
      pid: 999999999,
      acquiredAgoMs: 60_000,
      refreshedAgoMs: 0,
      command: '/path/to/gbrain/src/cli.ts serve',
      subcommand: 'serve',
    });

    const lock = await acquireLock(TEST_DIR, { timeoutMs: 2_000 });
    expect(lock.acquired).toBe(true);
    await releaseLock(lock);
  });

  test('[REGRESSION] a LIVE holder with a fresh heartbeat is NOT stolen even when the lock is old', async () => {
    // The WAL-corruption bug: a >5min embed used to get its lock force-removed.
    // Now an alive holder that heartbeated recently is left alone regardless of
    // age. acquired 20min ago, but refreshed just now → must wait, not steal.
    writeHolder({ pid: process.pid, acquiredAgoMs: 20 * 60_000, refreshedAgoMs: 0 });

    await expect(acquireLock(TEST_DIR, { timeoutMs: 1200 })).rejects.toThrow(/Timed out/);
    // Holder's lock still present (was never stolen).
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(true);
  });

  test('[REGRESSION #2348] a LIVE PID with a STALE heartbeat is NOT stolen', async () => {
    // The #2348 corruption: a live `gbrain dream`/embed holder whose heartbeat
    // lapsed (the JS event loop is blocked during a long synchronous WASM
    // import) used to get its lock reaped past the grace window — letting a
    // second OS process open the same data dir and corrupt the catalog +
    // pgvector extension state. A live PID is now NEVER stolen, regardless of
    // how stale its heartbeat is. Acquire must time out, not steal.
    writeHolder({ pid: process.pid, acquiredAgoMs: 25 * 60_000, refreshedAgoMs: 20 * 60_000 });

    await expect(acquireLock(TEST_DIR, { timeoutMs: 1200 })).rejects.toThrow(/Timed out/);
    // The live holder's lock is still present — never force-removed.
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(true);
  });

  test('explains live gbrain serve contention is not a sync advisory lock', async () => {
    writeHolder({
      pid: process.pid,
      acquiredAgoMs: 60_000,
      refreshedAgoMs: 0,
      command: 'bun /Users/master/.bun/bin/gbrain serve',
    });

    let message = '';
    try {
      await acquireLock(TEST_DIR, { timeoutMs: 100 });
    } catch (error) {
      message = error instanceof Error ? error.message : String(error);
    }
    expect(message).toContain('serve↔sync contention');
    expect(message).toContain('not the `gbrain-sync:*` advisory lock');
    expect(message).toContain('`gbrain sync --break-lock` will not clear a live PGLite holder');
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(true);
  });

  test('[REGRESSION] releaseLock does NOT remove a lock that was stolen + re-acquired by another process', async () => {
    // We acquire, then simulate a steal: another process reaped us past grace
    // and now owns the lock (different pid + acquired_at). Our releaseLock must
    // NOT delete their live lock — doing so would let a third process in
    // alongside the new owner (the #2058 corruption class).
    const lock: LockHandle = await acquireLock(TEST_DIR);
    expect(lock.acquired).toBe(true);
    expect(lock.ownerToken).toBeDefined();
    if (lock.heartbeat) clearInterval(lock.heartbeat); // stop our heartbeat for a deterministic test

    // Overwrite the lock file as if process B re-acquired it.
    const lockFile = join(TEST_DIR, '.gbrain-lock', 'lock');
    const bNow = Date.now() + 1;
    writeFileSync(lockFile, JSON.stringify({ pid: 999999, acquired_at: bNow, refreshed_at: bNow, command: 'process B' }));

    await releaseLock(lock); // our (stale) handle

    // B's lock survives — we did not clobber it.
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(true);
    const after = JSON.parse(readFileSync(lockFile, 'utf-8'));
    expect(after.pid).toBe(999999);

    // Cleanup for afterEach.
    rmSync(join(TEST_DIR, '.gbrain-lock'), { recursive: true, force: true });
  });

  test('acquire starts a heartbeat and seeds refreshed_at; release clears it', async () => {
    const lock: LockHandle = await acquireLock(TEST_DIR);
    expect(lock.acquired).toBe(true);
    expect(lock.heartbeat).toBeDefined();
    const data = JSON.parse(readFileSync(join(TEST_DIR, '.gbrain-lock', 'lock'), 'utf-8'));
    expect(data.refreshed_at).toBeDefined();
    expect(typeof data.refreshed_at).toBe('number');

    await releaseLock(lock);
    expect(lock.heartbeat).toBeUndefined();
    expect(existsSync(join(TEST_DIR, '.gbrain-lock'))).toBe(false);
  });
});

describe('pglite-lock reap classification (WAL-repair wave)', () => {
  // Unique per-test tmpdirs: the reap marker lands at `${dataDir}.lock-reap.json`
  // — a SIBLING of the data dir — so each test gets its own parent to rm.
  function freshDataDir(): { parent: string; dataDir: string } {
    const parent = mkdtempSync(join(tmpdir(), 'gbrain-lock-reap-'));
    return { parent, dataDir: join(parent, 'data') };
  }

  /**
   * A PID that provably belongs to no live process: spawn a short-lived child,
   * wait for it (spawnSync reaps it), then verify kill(pid, 0) throws. Retries
   * to dodge instant PID reuse.
   */
  function deadPid(): number {
    for (let attempt = 0; attempt < 5; attempt++) {
      const proc = Bun.spawnSync(['bash', '-c', 'exit 0']);
      const pid = proc.pid;
      try {
        process.kill(pid, 0); // still alive/visible → PID reused, try again
      } catch {
        return pid;
      }
    }
    throw new Error('could not obtain a provably-dead PID after 5 spawns');
  }

  test('corrupt lock file: reaped acquisition + persisted .lock-reap.json marker', async () => {
    const { parent, dataDir } = freshDataDir();
    try {
      const lockDir = join(dataDir, '.gbrain-lock');
      mkdirSync(lockDir, { recursive: true });
      writeFileSync(join(lockDir, 'lock'), 'not json {{{'); // holder liveness UNKNOWABLE

      const lock = await acquireLock(dataDir, { timeoutMs: 5000 });
      try {
        expect(lock.acquired).toBe(true);
        expect(lock.reaped).toBe(true);
        // Unknowable-liveness reap is persisted cross-process for the repair gate.
        expect(existsSync(`${dataDir}.lock-reap.json`)).toBe(true);
        const marker = JSON.parse(readFileSync(`${dataDir}.lock-reap.json`, 'utf-8'));
        expect(typeof marker.ts).toBe('number');
        expect(marker.by).toBe(process.pid);
      } finally {
        await releaseLock(lock);
      }
    } finally {
      rmSync(parent, { recursive: true, force: true });
    }
  });

  test('clean acquisition: reaped falsy, no .lock-reap.json marker', async () => {
    const { parent, dataDir } = freshDataDir();
    try {
      const lock = await acquireLock(dataDir, { timeoutMs: 5000 });
      try {
        expect(lock.acquired).toBe(true);
        expect(lock.reaped).toBeFalsy();
        expect(existsSync(`${dataDir}.lock-reap.json`)).toBe(false);
      } finally {
        await releaseLock(lock);
      }
    } finally {
      rmSync(parent, { recursive: true, force: true });
    }
  });

  test('dead-PID lock: reaped acquisition but NO marker (affirmative ESRCH verdict)', async () => {
    const { parent, dataDir } = freshDataDir();
    try {
      const lockDir = join(dataDir, '.gbrain-lock');
      mkdirSync(lockDir, { recursive: true });
      const now = Date.now();
      writeFileSync(join(lockDir, 'lock'), JSON.stringify({
        pid: deadPid(),
        acquired_at: now - 60_000,
        refreshed_at: now - 60_000,
        command: 'gbrain embed',
        subcommand: 'embed',
      }));

      const lock = await acquireLock(dataDir, { timeoutMs: 5000 });
      try {
        expect(lock.acquired).toBe(true);
        expect(lock.reaped).toBe(true);
        // Dead-PID reaps deliberately do NOT quarantine the next acquirer.
        expect(existsSync(`${dataDir}.lock-reap.json`)).toBe(false);
      } finally {
        await releaseLock(lock);
      }
    } finally {
      rmSync(parent, { recursive: true, force: true });
    }
  }, 30_000);
});