pgserve 2.3.0 → 2.5.0

package/src/cluster.js

@@ -1,654 +0,0 @@
-/**
- * Cluster Mode for pgserve
- *
- * Architecture:
- * - PRIMARY process: Runs single embedded PostgreSQL instance
- * - WORKER processes: Only run TCP routing to PRIMARY's PostgreSQL
- *
- * This enables multi-core scaling (3-5x throughput on multi-core systems)
- * while maintaining a single PostgreSQL instance.
- */
-
-import cluster from 'cluster';
-import os from 'os';
-import { SQL } from 'bun';
-import { createLogger } from './logger.js';
-import { PostgresManager } from './postgres.js';
-import { extractDatabaseName } from './protocol.js';
-import { EventEmitter } from 'events';
-import { loadEffectiveConfig } from './settings-loader.cjs';
-
-// PostgreSQL protocol constants
-const PROTOCOL_VERSION_3 = 196608;
-const SSL_REQUEST_CODE = 80877103;
-const GSSAPI_REQUEST_CODE = 80877104;
-const CANCEL_REQUEST_CODE = 80877102;
-
-/**
- * Attempt to write a pending buffer to a target socket.
- * Returns remaining unwritten bytes, or null if fully flushed.
- */
-function flushPending(target, pending) {
-  const written = target.write(pending);
-  if (written === pending.byteLength) return null;
-  if (written === 0) return pending;
-  return pending.subarray(written);
-}
-
-// Stats collection constants
-const WORKER_STATS_TIMEOUT_MS = 10000; // Worker stats older than this are considered stale
-const WORKER_STATS_REPORT_INTERVAL_MS = 4000; // How often workers report stats to primary
-
-/**
- * ClusterRouter - Lightweight TCP router for worker processes
- * Does NOT start PostgreSQL - connects to PRIMARY's PostgreSQL via Unix socket
- */
-class ClusterRouter extends EventEmitter {
-  constructor(options = {}) {
-    super();
-    this.port = options.port || 8432;
-    this.host = options.host || '127.0.0.1';
-    this.pgSocketPath = options.pgSocketPath; // From PRIMARY
-    this.pgPort = options.pgPort;
-    this.pgUser = options.pgUser || 'postgres';
-    this.pgPassword = options.pgPassword || 'postgres';
-    this.autoProvision = options.autoProvision !== false;
-    this.maxConnections = options.maxConnections || 1000;
-    this.enablePgvector = options.enablePgvector || false;
-
-    this.logger = createLogger({ level: options.logLevel || 'info' });
-    this.sql = null;  // Bun.sql for admin queries
-    this.server = null;
-    this.connections = new Set();
-    this.setMaxListeners(this.maxConnections + 10);
-
-    // Connection stats tracking for IPC reporting
-    this.connectionStats = {
-      totalConnected: 0,
-      totalDisconnected: 0
-    };
-  }
-
-  /**
-   * Socket state storage for Bun TCP handler model
-   */
-  socketState = new WeakMap();
-
-  async start() {
-    // Admin connection for auto-provisioning databases (Bun.sql)
-    if (this.autoProvision) {
-      // Bun.sql uses TCP connections - Unix sockets not directly supported
-      // This is fine for admin queries (low volume, local connection)
-      this.sql = new SQL({
-        hostname: '127.0.0.1',
-        port: this.pgPort,
-        database: 'postgres',
-        username: this.pgUser,
-        password: this.pgPassword,
-        max: 2,  // Small pool for admin queries
-        idleTimeout: 30,
-      });
-    }
-
-    // Create TCP server using Bun.listen() for 2-3x throughput
-    const router = this;
-    const isWindows = os.platform() === 'win32';
-    this.server = Bun.listen({
-      hostname: this.host,
-      port: this.port,
-      reusePort: !isWindows,  // SO_REUSEPORT for multi-worker port sharing (Linux/macOS only)
-      socket: {
-        data(socket, data) {
-          router.handleSocketData(socket, data);
-        },
-        open(socket) {
-          router.handleSocketOpen(socket);
-        },
-        close(socket) {
-          router.handleSocketClose(socket);
-        },
-        error(socket, error) {
-          router.handleSocketError(socket, error);
-        },
-        drain(socket) {
-          const state = router.socketState.get(socket);
-          if (!state) return;
-          // Flush any pending PG→Client data
-          if (state.pendingToClient) {
-            state.pendingToClient = flushPending(socket, state.pendingToClient);
-          }
-          // If fully flushed, resume reading from PostgreSQL
-          if (!state.pendingToClient && state.pgSocket) {
-            state.pgSocket.resume();
-          }
-        }
-      }
-    });
-
-    // Verify port actually bound (detect silent failures on Windows)
-    if (!this.server || !this.server.port) {
-      throw new Error(`Failed to bind to port ${this.port} - reusePort may not be supported on this platform`);
-    }
-
-    this.emit('listening');
-  }
-
-  async createDatabase(dbName) {
-    if (!this.autoProvision || !this.sql) return;
-
-    try {
-      // Bun.sql uses tagged template literals for parameterized queries
-      const result = await this.sql`SELECT 1 FROM pg_database WHERE datname = ${dbName}`;
-
-      if (result.length === 0) {
-        // Use sql() helper for safe identifier escaping (like CREATE DATABASE)
-        await this.sql.unsafe(`CREATE DATABASE "${dbName.replace(/"/g, '""')}"`);
-
-        // Auto-enable pgvector extension if configured
-        if (this.enablePgvector) {
-          await this.enablePgvectorExtension(dbName);
-        }
-      }
-    } catch (error) {
-      // Ignore "already exists" (race condition between workers)
-      if (!error.message?.includes('already exists')) {
-        this.logger.error({ database: dbName, err: error }, 'Failed to create database');
-      }
-    }
-  }
-
-  /**
-   * Enable pgvector extension on a database
-   * Creates a temporary connection to the specific database to run CREATE EXTENSION
-   * @param {string} dbName - Database name to enable pgvector on
-   */
-  async enablePgvectorExtension(dbName) {
-    let dbPool = null;
-
-    try {
-      // Create temporary connection to the specific database
-      dbPool = new SQL({
-        hostname: '127.0.0.1',
-        port: this.pgPort,
-        database: dbName,
-        username: this.pgUser,
-        password: this.pgPassword,
-        max: 1,
-        idleTimeout: 5,
-        connectionTimeout: 5,
-      });
-
-      // Enable pgvector extension
-      await dbPool.unsafe('CREATE EXTENSION IF NOT EXISTS vector');
-      this.logger.info({ dbName }, 'pgvector extension enabled');
-    } catch (error) {
-      // Log but don't fail database creation - pgvector might not be available
-      this.logger.warn({ dbName, err: error.message }, 'Failed to enable pgvector extension (non-fatal)');
-    } finally {
-      // Always close the temporary connection
-      if (dbPool) {
-        await dbPool.close().catch(() => {});
-      }
-    }
-  }
-
-  /**
-   * Handle socket open (Bun TCP handler)
-   */
-  handleSocketOpen(socket) {
-    this.socketState.set(socket, {
-      buffer: null,
-      pgSocket: null,
-      dbName: null,
-      handshakeComplete: false,
-      pendingToPg: null,
-      pendingToClient: null
-    });
-    this.connections.add(socket);
-    this.connectionStats.totalConnected++;
-  }
-
-  /**
-   * Handle socket data (Bun TCP handler)
-   */
-  handleSocketData(socket, data) {
-    const state = this.socketState.get(socket);
-    if (!state) return;
-
-    // If handshake complete, forward to PostgreSQL
-    if (state.handshakeComplete && state.pgSocket) {
-      // If there's already pending data, append to it
-      if (state.pendingToPg) {
-        state.pendingToPg = Buffer.concat([state.pendingToPg, data]);
-        return;
-      }
-      const written = state.pgSocket.write(data);
-      if (written < data.byteLength) {
-        // Partial write — buffer remainder and pause client
-        state.pendingToPg = written === 0 ? Buffer.from(data) : Buffer.from(data.subarray(written));
-        socket.pause();
-      }
-      return;
-    }
-
-    // Buffer data for startup message parsing
-    if (state.buffer) {
-      state.buffer = Buffer.concat([state.buffer, data]);
-    } else {
-      state.buffer = Buffer.from(data);
-    }
-
-    this.processStartupMessage(socket, state);
-  }
-
-  /**
-   * Process PostgreSQL startup message and establish proxy connection
-   */
-  async processStartupMessage(socket, state) {
-    const buffer = state.buffer;
-    if (!buffer || buffer.length < 8) return;
-
-    const messageLength = buffer.readUInt32BE(0);
-    if (buffer.length < messageLength) return;
-
-    const code = buffer.readUInt32BE(4);
-
-    // Handle SSL/GSSAPI/Cancel requests
-    if (code === SSL_REQUEST_CODE || code === GSSAPI_REQUEST_CODE) {
-      socket.write(Buffer.from('N'));
-      state.buffer = buffer.length > messageLength ? buffer.subarray(messageLength) : null;
-      if (state.buffer) await this.processStartupMessage(socket, state);
-      return;
-    }
-
-    if (code === CANCEL_REQUEST_CODE) {
-      socket.end();
-      return;
-    }
-
-    if (code !== PROTOCOL_VERSION_3) {
-      this.logger.warn({ code }, 'Unsupported protocol version');
-      socket.end();
-      return;
-    }
-
-    const startupMessage = buffer.subarray(0, messageLength);
-    const dbName = extractDatabaseName(startupMessage);
-    state.dbName = dbName;
-
-    try {
-      await this.createDatabase(dbName);
-
-      const router = this;
-
-      // Shared handler for pgSocket (used by both unix and TCP paths)
-      const pgHandler = {
-        data(_pgSocket, pgData) {
-          // Forward PostgreSQL response to client with backpressure
-          if (state.pendingToClient) {
-            state.pendingToClient = Buffer.concat([state.pendingToClient, pgData]);
-            return;
-          }
-          const written = socket.write(pgData);
-          if (written < pgData.byteLength) {
-            state.pendingToClient = written === 0 ? Buffer.from(pgData) : Buffer.from(pgData.subarray(written));
-            _pgSocket.pause();
-          }
-        },
-        open(pgSocket) {
-          pgSocket.write(startupMessage);
-          state.handshakeComplete = true;
-        },
-        close(_pgSocket) {
-          socket.end();
-        },
-        error(_pgSocket, error) {
-          router.logger.error({ dbName, err: error }, 'PostgreSQL socket error');
-          socket.end();
-        },
-        drain(_pgSocket) {
-          // Flush any pending Client→PG data
-          if (state.pendingToPg) {
-            state.pendingToPg = flushPending(_pgSocket, state.pendingToPg);
-          }
-          // If fully flushed, resume reading from client
-          if (!state.pendingToPg) {
-            socket.resume();
-          }
-        }
-      };
-
-      if (this.pgSocketPath) {
-        state.pgSocket = await Bun.connect({ unix: this.pgSocketPath, socket: pgHandler });
-      } else {
-        state.pgSocket = await Bun.connect({ hostname: '127.0.0.1', port: this.pgPort, socket: pgHandler });
-      }
-    } catch (error) {
-      this.logger.error({ dbName, err: error }, 'Connection error');
-      socket.end();
-    }
-  }
-
-  /**
-   * Handle socket close (Bun TCP handler)
-   */
-  handleSocketClose(socket) {
-    const state = this.socketState.get(socket);
-    if (state) {
-      state.pendingToPg = null;
-      state.pendingToClient = null;
-      if (state.pgSocket) state.pgSocket.end();
-    }
-    this.connections.delete(socket);
-    this.socketState.delete(socket);
-    this.connectionStats.totalDisconnected++;
-  }
-
-  /**
-   * Get router stats for IPC reporting
-   */
-  getStats() {
-    return {
-      connections: this.connections.size,
-      totalConnected: this.connectionStats.totalConnected,
-      totalDisconnected: this.connectionStats.totalDisconnected,
-      pid: process.pid
-    };
-  }
-
-  /**
-   * Handle socket error (Bun TCP handler)
-   */
-  handleSocketError(socket, error) {
-    const state = this.socketState.get(socket);
-    if (error.code !== 'ECONNRESET') {
-      this.logger.error({ err: error, dbName: state?.dbName }, 'Socket error');
-    }
-    if (state) {
-      state.pendingToPg = null;
-      state.pendingToClient = null;
-      if (state.pgSocket) state.pgSocket.end();
-    }
-    this.connections.delete(socket);
-    this.socketState.delete(socket);
-  }
-
-  async stop() {
-    for (const socket of this.connections) {
-      socket.end();
-    }
-    this.connections.clear();
-
-    if (this.sql) {
-      try {
-        await this.sql.close();
-      } catch {
-        // Expected: connection may already be terminated during cleanup
-      }
-    }
-
-    // Close TCP server (Bun.listen returns a server with stop() method)
-    if (this.server) {
-      this.server.stop();
-    }
-  }
-}
-
-
-/**
- * Build a `backendExited` handler for cluster mode supervision.
- *
- * On unexpected exit (`expected: false`) — postgres SIGKILL'd, OOM-killed,
- * segfaulted, etc. — the handler:
- *   1. logs the exit code,
- *   2. flips `shuttingDown` so the cluster.on('exit') worker-respawn path
- *      no longer forks new workers,
- *   3. SIGTERMs every live worker (no point routing to a dead backend), and
- *   4. calls `exitFn(1)` so the parent process supervisor restarts us.
- *
- * On clean exit (`expected: true`, initiated by `pgManager.stop()`) the
- * handler is silent — the surrounding shutdown logic handles teardown.
- *
- * Exported so the supervision contract can be unit-tested without spawning
- * a real cluster (the integration path is already covered for single-process
- * mode by tests/wrapper-supervision.test.js).
- *
- * @param {object} args
- * @param {Map<number, {kill: (sig: string) => void}>} args.workers - live worker registry
- * @param {(v: boolean) => void} args.setShuttingDown - flips outer-scope `shuttingDown`
- * @param {(code: number) => void} [args.exitFn=process.exit] - test seam
- * @param {(...args: unknown[]) => void} [args.log=console.error] - test seam
- * @returns {(info: {code: number, expected: boolean}) => void}
- */
-export function buildClusterSupervisionHandler({
-  workers,
-  setShuttingDown,
-  exitFn = process.exit,
-  log = console.error,
-}) {
-  return ({ code, expected }) => {
-    if (expected) return;
-    log(
-      `[pgserve] postgres backend exited unexpectedly (code=${code}) in cluster mode; ` +
-      `the primary is exiting so a process supervisor can restart it.`
-    );
-    setShuttingDown(true);
-    for (const worker of workers.values()) {
-      try { worker.kill('SIGTERM'); } catch { /* worker may already be dead */ }
-    }
-    exitFn(1);
-  };
-}
-
-/**
- * Start pgserve in cluster mode
- */
-export async function startClusterServer(options = {}) {
-  const numWorkers = options.workers || os.cpus().length;
-  const port = options.port || 8432;
-  const host = options.host || '127.0.0.1';
-  const pgPort = options.pgPort || (port + 1000);
-
-  if (cluster.isPrimary) {
-    // Port binding happens in workers via Bun.listen with reusePort
-    // If port is in use, first worker will fail with EADDRINUSE
-    console.log(`[pgserve] Cluster mode: ${numWorkers} workers`);
-
-    // PRIMARY: Start our embedded PostgreSQL (single instance)
-    const logger = createLogger({ level: options.logLevel || 'info' });
-    const pgManager = new PostgresManager({
-      dataDir: options.baseDir,
-      port: pgPort,
-      logger: logger.child({ component: 'postgres' }),
-      useRam: options.useRam,  // Use /dev/shm for true RAM storage (Linux only)
-      enablePgvector: options.enablePgvector  // Auto-enable pgvector extension on new databases
-    });
-
-    await pgManager.start();
-    const pgSocketPath = pgManager.getSocketPath();
-
-    console.log(`[pgserve] Embedded PostgreSQL started`);
-    console.log(`[pgserve] Socket: ${pgSocketPath || `TCP port ${pgPort}`}`);
-
-    // Track shutdown state and worker registry early so the supervision
-    // handler below can tear workers down on unexpected backend death.
-    let shuttingDown = false;
-    const workers = new Map();
-    const workerStats = new Map(); // Track stats from each worker
-
-    // Supervision: when the embedded postgres backend dies unexpectedly
-    // (SIGKILL/OOM/segfault — anything other than a clean stop()), exit the
-    // primary so a process supervisor (`genie serve`, pm2, systemd) restarts
-    // the cluster cleanly with a fresh backend. Without this, the primary
-    // keeps running with a zombie pgManager (socketDir nulled) and every
-    // worker fails StartupMessages with "Connection closed" forever while
-    // pm2 reports the process as healthy. Mirrors the single-process fix
-    // in bin/postgres-server.js (PgserveDaemon.on('backendDiedUnexpectedly'))
-    // — pgserve#45 only protected the daemon path, not cluster mode (default
-    // on multi-core systems).
-    pgManager.on('backendExited', buildClusterSupervisionHandler({
-      workers,
-      setShuttingDown: (v) => { shuttingDown = v; },
-    }));
-
-    // Fork workers with PostgreSQL connection info.
-    //
-    // Pass through using AUTOPG_<X> (the primary names) so workers don't
-    // emit the legacy-PGSERVE deprecation log for our own internal IPC.
-    // PGSERVE_WORKER stays as-is — it's an internal flag, not part of the
-    // settings precedence chain.
-    const workerEnv = () => ({
-      PGSERVE_WORKER: 'true',
-      AUTOPG_PORT: String(port),
-      AUTOPG_HOST: host,
-      AUTOPG_PG_SOCKET: pgSocketPath || '',
-      AUTOPG_PG_PORT: String(pgPort),
-      AUTOPG_PG_USER: 'postgres',
-      AUTOPG_PG_PASSWORD: 'postgres',
-      AUTOPG_LOG_LEVEL: options.logLevel || 'info',
-      AUTOPG_AUTO_PROVISION: options.autoProvision !== false ? 'true' : 'false',
-      // max_connections is a postgres GUC, not a server-level env. Workers
-      // read it from settings.postgres via loadEffectiveConfig; we still
-      // ship the value here so a CLI override (e.g. `--max-connections`)
-      // reaches the worker before the daemon writes settings.json.
-      AUTOPG_MAX_CONNECTIONS: String(options.maxConnections || 1000),
-      AUTOPG_ENABLE_PGVECTOR: options.enablePgvector ? 'true' : 'false',
-    });
-    for (let i = 0; i < numWorkers; i++) {
-      const worker = cluster.fork(workerEnv());
-      workers.set(worker.id, worker);
-    }
-
-    // (shuttingDown declared above with the supervision handler.)
-
-    // Restart dead workers (unless shutting down)
-    cluster.on('exit', (worker, code, signal) => {
-      workers.delete(worker.id);
-
-      if (shuttingDown) {
-        return; // Don't restart during shutdown
-      }
-
-      console.log(`[pgserve] Worker ${worker.id} died (${signal || code}), restarting...`);
-      const newWorker = cluster.fork(workerEnv());
-      workers.set(newWorker.id, newWorker);
-    });
-
-    // Wait for workers to be ready and handle IPC messages
-    let readyCount = 0;
-    await new Promise((resolve) => {
-      cluster.on('message', (worker, message) => {
-        if (message.type === 'ready') {
-          readyCount++;
-          if (readyCount === numWorkers) resolve();
-        } else if (message.type === 'stats') {
-          // Update worker stats from IPC
-          workerStats.set(worker.id, {
-            ...message.data,
-            lastUpdate: Date.now()
-          });
-        }
-      });
-    });
-
-    console.log(`[pgserve] All ${numWorkers} workers ready`);
-    console.log(`[pgserve] Listening on ${host}:${port}`);
-
-    return {
-      workers,
-      pgPort,
-      pgSocketPath,
-      stop: async () => {
-        console.log('[pgserve] Stopping cluster...');
-        shuttingDown = true; // Prevent worker restart during shutdown
-        for (const worker of workers.values()) {
-          worker.send({ type: 'shutdown' });
-        }
-        await new Promise((resolve) => {
-          const check = setInterval(() => {
-            if (workers.size === 0) {
-              clearInterval(check);
-              resolve();
-            }
-          }, 100);
-        });
-        await pgManager.stop();
-        console.log('[pgserve] Cluster stopped');
-      },
-      getStats: () => {
-        // Aggregate stats from all workers
-        let totalConnections = 0;
-        let totalConnected = 0;
-        let totalDisconnected = 0;
-        const activeWorkerStats = {};
-
-        for (const [id, stats] of workerStats) {
-          // Only include recent stats (within timeout window)
-          if (Date.now() - stats.lastUpdate < WORKER_STATS_TIMEOUT_MS) {
-            totalConnections += stats.connections || 0;
-            totalConnected += stats.totalConnected || 0;
-            totalDisconnected += stats.totalDisconnected || 0;
-            activeWorkerStats[id] = stats;
-          }
-        }
-
-        return {
-          workers: workers.size,
-          pids: Array.from(workers.values()).map(w => w.process.pid),
-          connections: {
-            active: totalConnections,
-            totalConnected,
-            totalDisconnected
-          },
-          workerStats: activeWorkerStats
-        };
-      },
-      pgManager
-    };
-  } else {
-    // WORKER: Only run TCP routing, connect to PRIMARY's PostgreSQL.
-    //
-    // Worker config comes from loadEffectiveConfig() (defaults < file < env).
-    // The primary fork() above sets PGSERVE_* env vars so existing supervised
-    // installs keep working; AUTOPG_* env vars take precedence when set, and
-    // a one-time deprecation note is emitted for legacy-only PGSERVE_* hits.
-    const { settings } = loadEffectiveConfig();
-    const router = new ClusterRouter({
-      port: settings.server.port,
-      host: settings.server.host,
-      pgSocketPath: settings.server.pgSocketPath || null,
-      pgPort: settings.server.pgPort,
-      pgUser: settings.server.pgUser,
-      pgPassword: settings.server.pgPassword,
-      logLevel: settings.runtime.logLevel,
-      autoProvision: settings.runtime.autoProvision,
-      maxConnections: settings.postgres.max_connections,
-      enablePgvector: settings.runtime.enablePgvector,
-    });
-
-    await router.start();
-
-    // Tell PRIMARY we're ready
-    process.send({ type: 'ready' });
-
-    // Periodically send stats to PRIMARY
-    const statsInterval = setInterval(() => {
-      try {
-        process.send({ type: 'stats', data: router.getStats() });
-      } catch {
-        // Expected: IPC channel may be closed during shutdown
-      }
-    }, WORKER_STATS_REPORT_INTERVAL_MS);
-
-    // Handle shutdown
-    process.on('message', async (message) => {
-      if (message.type === 'shutdown') {
-        clearInterval(statsInterval);
-        await router.stop();
-        process.exit(0);
-      }
-    });
-
-    return router;
-  }
-}