@happier-dev/stack 0.2.1-preview.1775586714.26562 → 0.2.11

package/scripts/utils/net/tcp_forward.mjs

@@ -12,6 +12,8 @@
 
 import net from 'node:net';
 
+const stopPromises = new WeakMap();
+
 /**
  * Create a TCP forwarding server.
  *
@@ -24,12 +26,24 @@ import net from 'node:net';
  * @returns {net.Server}
  */
 export function createTcpForwarder({ listenHost, listenPort, targetHost, targetPort, label = 'tcp-forward' }) {
+  let currentTarget = normalizeUpstream({ targetHost, targetPort });
+  const activeConnections = new Set();
+
   const server = net.createServer((clientSocket) => {
-    const targetSocket = net.createConnection({ host: targetHost, port: targetPort }, () => {
+    const target = currentTarget;
+    const targetSocket = net.createConnection({ host: target.targetHost, port: target.targetPort }, () => {
       // Connection established, pipe data both ways
       clientSocket.pipe(targetSocket);
       targetSocket.pipe(clientSocket);
     });
+    const connection = { clientSocket, targetSocket, target };
+    activeConnections.add(connection);
+
+    const cleanup = () => {
+      activeConnections.delete(connection);
+      clientSocket.unpipe(targetSocket);
+      targetSocket.unpipe(clientSocket);
+    };
 
     // Handle errors on both sockets
     clientSocket.on('error', (err) => {
@@ -47,10 +61,54 @@ export function createTcpForwarder({ listenHost, listenPort, targetHost, targetP
     });
 
     // Clean up on close
-    clientSocket.on('close', () => targetSocket.destroy());
-    targetSocket.on('close', () => clientSocket.destroy());
+    clientSocket.on('close', () => {
+      cleanup();
+      targetSocket.destroy();
+    });
+    targetSocket.on('close', () => {
+      cleanup();
+      clientSocket.destroy();
+    });
   });
 
+  server.setUpstream = (nextTarget) => {
+    currentTarget = normalizeUpstream(nextTarget);
+    return { ...currentTarget };
+  };
+
+  server.getUpstream = () => ({ ...currentTarget });
+
+  server.getActiveConnectionCount = () => activeConnections.size;
+
+  server.closeConnectionsAfterGrace = async ({ graceMs = 0, target = null } = {}) => {
+    const delayMs = Math.max(0, Number(graceMs) || 0);
+    if (delayMs > 0) {
+      await new Promise((resolve) => setTimeout(resolve, delayMs));
+    }
+
+    const normalizedTarget = target ? normalizeUpstream(target) : null;
+    const connectionsToClose = [...activeConnections].filter((connection) => (
+      !normalizedTarget || upstreamsEqual(connection.target, normalizedTarget)
+    ));
+    const closeWaits = connectionsToClose.map(({ clientSocket, targetSocket }) => Promise.all([
+      waitForSocketClose(clientSocket),
+      waitForSocketClose(targetSocket),
+    ]));
+
+    for (const connection of connectionsToClose) {
+      connection.clientSocket.destroy();
+      connection.targetSocket.destroy();
+    }
+
+    if (connectionsToClose.length === 0) return;
+
+    await Promise.race([
+      Promise.all(closeWaits),
+      new Promise((resolve) => setTimeout(resolve, 250)),
+    ]);
+  };
+  server.closeIdleOrAllConnectionsAfterGrace = server.closeConnectionsAfterGrace;
+
   server.on('error', (err) => {
     process.stderr.write(`[${label}] server error: ${err.message}\n`);
   });
@@ -58,6 +116,25 @@ export function createTcpForwarder({ listenHost, listenPort, targetHost, targetP
   return server;
 }
 
+function normalizeUpstream({ targetHost, targetPort }) {
+  const host = String(targetHost || '').trim();
+  const port = Number(targetPort);
+  if (!host) throw new Error('targetHost is required');
+  if (!Number.isInteger(port) || port <= 0 || port > 65535) {
+    throw new Error(`targetPort must be a valid TCP port, got ${targetPort}`);
+  }
+  return { targetHost: host, targetPort: port };
+}
+
+function upstreamsEqual(left, right) {
+  return left.targetHost === right.targetHost && left.targetPort === right.targetPort;
+}
+
+function waitForSocketClose(socket) {
+  if (!socket || socket.closed) return Promise.resolve();
+  return new Promise((resolve) => socket.once('close', resolve));
+}
+
 /**
  * Start a TCP forwarder and return a promise that resolves when listening.
  *
@@ -100,16 +177,47 @@ function trySendIpc(msg) {
  */
 export async function stopTcpForwarder(server, label = 'tcp-forward') {
   if (!server) return;
-  return new Promise((resolve) => {
-    server.close(() => {
+  const existingStop = stopPromises.get(server);
+  if (existingStop) return existingStop;
+
+  const stopPromise = new Promise((resolve) => {
+    let done = false;
+    const finish = () => {
+      if (done) return;
+      done = true;
       process.stdout.write(`[${label}] stopped\n`);
       resolve();
+    };
+
+    const closeActive = typeof server.closeConnectionsAfterGrace === 'function'
+      ? server.closeConnectionsAfterGrace({ graceMs: 0 })
+      : typeof server.closeIdleOrAllConnectionsAfterGrace === 'function'
+        ? server.closeIdleOrAllConnectionsAfterGrace({ graceMs: 0 })
+      : Promise.resolve();
+
+    closeActive.then(() => {
+      if (!server.listening) finish();
+    }, () => {
+      finish();
     });
-    // Force-close after timeout
-    setTimeout(() => {
-      resolve();
-    }, 2000);
+
+    try {
+      server.close(() => {
+        closeActive.finally(finish);
+      });
+    } catch (err) {
+      if (err && typeof err === 'object' && err.code === 'ERR_SERVER_NOT_RUNNING') {
+        closeActive.finally(finish);
+        return;
+      }
+      throw err;
+    }
+
+    // Force-close after timeout.
+    setTimeout(finish, 2000);
   });
+  stopPromises.set(server, stopPromise);
+  return stopPromise;
 }
 
 // Standalone CLI mode

package/scripts/utils/net/tcp_forward.test.mjs

@@ -0,0 +1,255 @@
+import test from 'node:test';
+import assert from 'node:assert/strict';
+import net from 'node:net';
+
+import { startTcpForwarder, stopTcpForwarder } from './tcp_forward.mjs';
+
+async function startEchoServer(label) {
+  const sockets = new Set();
+  const server = net.createServer((socket) => {
+    sockets.add(socket);
+    socket.on('data', (chunk) => {
+      socket.write(`${label}:${chunk.toString()}`);
+    });
+    socket.on('close', () => sockets.delete(socket));
+  });
+
+  await new Promise((resolve, reject) => {
+    server.once('error', reject);
+    server.listen({ host: '127.0.0.1', port: 0 }, resolve);
+  });
+
+  const addr = server.address();
+  const port = typeof addr === 'object' && addr ? addr.port : null;
+  if (!port) throw new Error(`failed to start ${label} echo server`);
+
+  return {
+    port,
+    sockets,
+    async stop() {
+      for (const socket of sockets) socket.destroy();
+      await closeServer(server);
+    },
+  };
+}
+
+async function getUnusedLoopbackPort() {
+  const server = net.createServer();
+  await new Promise((resolve, reject) => {
+    server.once('error', reject);
+    server.listen({ host: '127.0.0.1', port: 0 }, resolve);
+  });
+  const addr = server.address();
+  const port = typeof addr === 'object' && addr ? addr.port : null;
+  await closeServer(server);
+  if (!port) throw new Error('failed to allocate an unused port');
+  return port;
+}
+
+async function closeServer(server) {
+  await new Promise((resolve) => {
+    server.close(() => resolve());
+  });
+}
+
+async function connectClient(port) {
+  const socket = net.createConnection({ host: '127.0.0.1', port });
+  await new Promise((resolve, reject) => {
+    socket.once('connect', resolve);
+    socket.once('error', reject);
+  });
+  return socket;
+}
+
+async function writeAndRead(socket, message) {
+  const dataPromise = new Promise((resolve, reject) => {
+    const onData = (chunk) => {
+      cleanup();
+      resolve(chunk.toString());
+    };
+    const onError = (err) => {
+      cleanup();
+      reject(err);
+    };
+    const cleanup = () => {
+      socket.off('data', onData);
+      socket.off('error', onError);
+    };
+    socket.once('data', onData);
+    socket.once('error', onError);
+  });
+  socket.write(message);
+  return dataPromise;
+}
+
+async function waitForSocketClosed(socket) {
+  if (socket.destroyed) return;
+  await new Promise((resolve) => {
+    socket.once('close', resolve);
+  });
+}
+
+test('upstream flip affects new connections while existing connections stay on their original upstream until drain', { timeout: 5000 }, async () => {
+  const oldUpstream = await startEchoServer('old');
+  const newUpstream = await startEchoServer('new');
+  let forwarder;
+  const sockets = [];
+  try {
+    forwarder = await startTcpForwarder({
+      listenHost: '127.0.0.1',
+      listenPort: 0,
+      targetHost: '127.0.0.1',
+      targetPort: oldUpstream.port,
+      label: 'tcp-forward-test',
+    });
+
+    assert.equal(typeof forwarder.server.setUpstream, 'function');
+    assert.equal(typeof forwarder.server.closeIdleOrAllConnectionsAfterGrace, 'function');
+
+    const existing = await connectClient(forwarder.port);
+    sockets.push(existing);
+    assert.equal(await writeAndRead(existing, 'one'), 'old:one');
+
+    forwarder.server.setUpstream({ targetHost: '127.0.0.1', targetPort: newUpstream.port });
+
+    const fresh = await connectClient(forwarder.port);
+    sockets.push(fresh);
+    assert.equal(await writeAndRead(fresh, 'two'), 'new:two');
+    assert.equal(await writeAndRead(existing, 'three'), 'old:three');
+
+    const drain = forwarder.server.closeIdleOrAllConnectionsAfterGrace({ graceMs: 10 });
+    await Promise.all([drain, waitForSocketClosed(existing), waitForSocketClosed(fresh)]);
+
+    const afterDrain = await connectClient(forwarder.port);
+    sockets.push(afterDrain);
+    assert.equal(await writeAndRead(afterDrain, 'four'), 'new:four');
+  } finally {
+    for (const socket of sockets) socket.destroy();
+    await stopTcpForwarder(forwarder?.server, 'tcp-forward-test');
+    await oldUpstream.stop();
+    await newUpstream.stop();
+  }
+});
+
+test('targeted drain closes only connections bound to the selected upstream', { timeout: 5000 }, async () => {
+  const oldUpstream = await startEchoServer('old');
+  const newUpstream = await startEchoServer('new');
+  let forwarder;
+  const sockets = [];
+  try {
+    forwarder = await startTcpForwarder({
+      listenHost: '127.0.0.1',
+      listenPort: 0,
+      targetHost: '127.0.0.1',
+      targetPort: oldUpstream.port,
+      label: 'tcp-forward-target-drain-test',
+    });
+
+    const oldConnection = await connectClient(forwarder.port);
+    sockets.push(oldConnection);
+    assert.equal(await writeAndRead(oldConnection, 'one'), 'old:one');
+
+    forwarder.server.setUpstream({ targetHost: '127.0.0.1', targetPort: newUpstream.port });
+
+    const newConnection = await connectClient(forwarder.port);
+    sockets.push(newConnection);
+    assert.equal(await writeAndRead(newConnection, 'two'), 'new:two');
+
+    await forwarder.server.closeConnectionsAfterGrace({
+      graceMs: 0,
+      target: { targetHost: '127.0.0.1', targetPort: oldUpstream.port },
+    });
+
+    await waitForSocketClosed(oldConnection);
+    assert.equal(await writeAndRead(newConnection, 'three'), 'new:three');
+  } finally {
+    for (const socket of sockets) socket.destroy();
+    await stopTcpForwarder(forwarder?.server, 'tcp-forward-target-drain-test');
+    await oldUpstream.stop();
+    await newUpstream.stop();
+  }
+});
+
+test('forwarder shutdown closes listener and active sockets', { timeout: 5000 }, async () => {
+  const upstream = await startEchoServer('upstream');
+  let forwarder;
+  let client;
+  try {
+    forwarder = await startTcpForwarder({
+      listenHost: '127.0.0.1',
+      listenPort: 0,
+      targetHost: '127.0.0.1',
+      targetPort: upstream.port,
+      label: 'tcp-forward-shutdown-test',
+    });
+
+    client = await connectClient(forwarder.port);
+    assert.equal(await writeAndRead(client, 'before'), 'upstream:before');
+
+    await stopTcpForwarder(forwarder.server, 'tcp-forward-shutdown-test');
+    await waitForSocketClosed(client);
+
+    await assert.rejects(
+      () => connectClient(forwarder.port),
+      (err) => err && err.code === 'ECONNREFUSED',
+    );
+  } finally {
+    client?.destroy();
+    await stopTcpForwarder(forwarder?.server, 'tcp-forward-shutdown-test');
+    await upstream.stop();
+  }
+});
+
+test('connection drain waits for socket close before resolving', { timeout: 5000 }, async () => {
+  const upstream = await startEchoServer('upstream');
+  let forwarder;
+  let client;
+  try {
+    forwarder = await startTcpForwarder({
+      listenHost: '127.0.0.1',
+      listenPort: 0,
+      targetHost: '127.0.0.1',
+      targetPort: upstream.port,
+      label: 'tcp-forward-drain-close-test',
+    });
+
+    client = await connectClient(forwarder.port);
+    assert.equal(await writeAndRead(client, 'before'), 'upstream:before');
+
+    await forwarder.server.closeIdleOrAllConnectionsAfterGrace({ graceMs: 0 });
+
+    assert.equal(forwarder.server.getActiveConnectionCount(), 0);
+  } finally {
+    client?.destroy();
+    await stopTcpForwarder(forwarder?.server, 'tcp-forward-drain-close-test');
+    await upstream.stop();
+  }
+});
+
+test('refused upstream closes the client without unhandled rejections', { timeout: 5000 }, async () => {
+  const targetPort = await getUnusedLoopbackPort();
+  const unhandled = [];
+  const onUnhandled = (reason) => unhandled.push(reason);
+  process.on('unhandledRejection', onUnhandled);
+
+  let forwarder;
+  let client;
+  try {
+    forwarder = await startTcpForwarder({
+      listenHost: '127.0.0.1',
+      listenPort: 0,
+      targetHost: '127.0.0.1',
+      targetPort,
+      label: 'tcp-forward-refused-test',
+    });
+
+    client = await connectClient(forwarder.port);
+    await waitForSocketClosed(client);
+    await new Promise((resolve) => setImmediate(resolve));
+    assert.deepEqual(unhandled, []);
+  } finally {
+    process.off('unhandledRejection', onUnhandled);
+    client?.destroy();
+    await stopTcpForwarder(forwarder?.server, 'tcp-forward-refused-test');
+  }
+});

package/scripts/utils/pglite_lock.mjs

@@ -1,11 +1,20 @@
 import { existsSync } from 'node:fs';
+import { createHash } from 'node:crypto';
 import { open, readFile, unlink } from 'node:fs/promises';
 import { dirname, join, resolve } from 'node:path';
 import { isPidAlive } from './proc/pids.mjs';
 import { getPidStartTime } from './proc/ownership.mjs';
 
-function lockPathForDbDir(dbDir) {
-  return join(dirname(dbDir), '.happier.pglite.lock');
+const PGLITE_LOCK_FILENAME = '.happier.pglite.lock';
+
+export function getPgliteDirLockPath(dbDir) {
+  const resolvedDbDir = resolve(String(dbDir ?? ''));
+  const short = createHash('sha256').update(resolvedDbDir).digest('hex').slice(0, 12);
+  return join(dirname(resolvedDbDir), `${PGLITE_LOCK_FILENAME}.${short}`);
+}
+
+function legacyAdjacentLockPathForDbDir(dbDir) {
+  return join(dirname(dbDir), PGLITE_LOCK_FILENAME);
 }
 
 function safeParseJson(raw) {
@@ -36,10 +45,43 @@ async function removeLock(lockPath) {
   }
 }
 
+async function assertNoLiveLegacyAdjacentLock(resolvedDbDir) {
+  const legacyLockPath = legacyAdjacentLockPathForDbDir(resolvedDbDir);
+  const hashedLockPath = getPgliteDirLockPath(resolvedDbDir);
+  if (legacyLockPath === hashedLockPath) return;
+
+  const existing = await readLock(legacyLockPath);
+  if (!existing) return;
+  if (existing.invalid) {
+    throw new Error(`pglite legacy lock is invalid: ${legacyLockPath}`);
+  }
+
+  const existingPid = Number(existing?.pid);
+  const existingDbDir = existing?.dbDir ? resolve(String(existing.dbDir)) : '';
+  const sameDbDir = existingDbDir ? existingDbDir === resolvedDbDir : true;
+  if (!sameDbDir) return;
+
+  if (Number.isFinite(existingPid) && isPidAlive(existingPid)) {
+    const meta = [
+      existing?.purpose ? `purpose=${existing.purpose}` : '',
+      existing?.createdAt ? `createdAt=${existing.createdAt}` : '',
+      existingDbDir ? `dbDir=${existingDbDir}` : '',
+    ]
+      .filter(Boolean)
+      .join(' ');
+    throw new Error(`pglite db dir is in use by pid=${existingPid}${meta ? ` (${meta})` : ''}`);
+  }
+
+  await removeLock(legacyLockPath);
+}
+
 export async function acquirePgliteDirLock(dbDir, { purpose = 'unknown' } = {}) {
   const resolvedDbDir = resolve(String(dbDir ?? ''));
   if (!resolvedDbDir) throw new Error('Missing dbDir');
-  const lockPath = lockPathForDbDir(resolvedDbDir);
+  await assertNoLiveLegacyAdjacentLock(resolvedDbDir);
+  const legacyInsideDbDir = join(resolvedDbDir, PGLITE_LOCK_FILENAME);
+  await removeLock(legacyInsideDbDir);
+  const lockPath = getPgliteDirLockPath(resolvedDbDir);
 
   for (let attempt = 0; attempt < 3; attempt += 1) {
     try {
@@ -105,3 +147,42 @@ export async function acquirePgliteDirLock(dbDir, { purpose = 'unknown' } = {})
 
   throw new Error('Failed to acquire pglite lock after retries');
 }
+
+export async function waitForPgliteDirLockRelease(
+  dbDir,
+  { timeoutMs = 5_000, intervalMs = 100, sleep = (ms) => new Promise((resolvePromise) => setTimeout(resolvePromise, ms)) } = {}
+) {
+  const resolvedDbDir = resolve(String(dbDir ?? ''));
+  if (!resolvedDbDir) return true;
+  const lockPaths = [getPgliteDirLockPath(resolvedDbDir), legacyAdjacentLockPathForDbDir(resolvedDbDir)];
+  const deadline = Date.now() + timeoutMs;
+
+  while (Date.now() < deadline) {
+    let held = false;
+    for (const lockPath of lockPaths) {
+      // eslint-disable-next-line no-await-in-loop
+      const existing = await readLock(lockPath);
+      if (!existing) continue;
+      if (existing.invalid) {
+        held = true;
+        continue;
+      }
+      const existingDbDir = existing?.dbDir ? resolve(String(existing.dbDir)) : '';
+      if (lockPath === legacyAdjacentLockPathForDbDir(resolvedDbDir) && existingDbDir && existingDbDir !== resolvedDbDir) {
+        continue;
+      }
+      const existingPid = Number(existing?.pid);
+      if (Number.isFinite(existingPid) && isPidAlive(existingPid)) {
+        held = true;
+        continue;
+      }
+      // eslint-disable-next-line no-await-in-loop
+      await removeLock(lockPath);
+    }
+    if (!held) return true;
+    // eslint-disable-next-line no-await-in-loop
+    await sleep(intervalMs);
+  }
+
+  return false;
+}