P1: Graceful shutdown — drain, flush traces, uncaught handlers (#170)

.env.example.compose

@@ -27,6 +27,7 @@ JAEGER=jaegertracing/all-in-one:latest
 
 # Fields for App (Required)
 PORT=8080
+SHUTDOWN_TIMEOUT_MS=20000
 LOG_LEVEL="info"
 CORS_ORIGIN="*"
 READINESS_PING_TIMEOUT_MS=2000

.env.example.lightnet

@@ -1,4 +1,5 @@
 PORT=8080
+SHUTDOWN_TIMEOUT_MS=20000
 LOG_LEVEL="info"
 CORS_ORIGIN="*"
 READINESS_PING_TIMEOUT_MS=2000

docs/getting-started.md

@@ -176,6 +176,7 @@ The server reads config from environment variables. `PG_CONN` is the only requir
 | --- | --- | --- |
 | `PG_CONN` | *(required)* | Postgres connection string for the archive-node DB |
 | `PORT` | `8080` | Port the GraphQL server listens on |
+| `SHUTDOWN_TIMEOUT_MS` | `20000` | Max ms to drain in-flight requests on SIGTERM before forcing exit. Keep Kubernetes `terminationGracePeriodSeconds` above this value |
 | `LOG_LEVEL` | `info` | `debug` \| `info` \| `warn` \| `error` |
 | `CORS_ORIGIN` | `*` | CORS allowed origin |
 | `READINESS_PING_TIMEOUT_MS` | `2000` | Upper bound on the `/readiness` database ping. Exceeding it returns 503 rather than leaving the probe to hang. Keep it below the orchestrator's probe `timeoutSeconds` |

src/envionment.d.ts

@@ -4,6 +4,7 @@ declare global {
       LOG_LEVEL: string;
       PORT?: string;
       PG_CONN: string;
+      SHUTDOWN_TIMEOUT_MS?: string;
       CORS_ORIGIN?: string;
       READINESS_PING_TIMEOUT_MS?: string;
       ENABLE_LOGGING?: bool;

src/index.ts

@@ -3,26 +3,80 @@
 import { buildContext } from './context.js';
 import { buildServer } from './server/server.js';
 import { buildPlugins } from './server/plugins.js';
+import { createGracefulShutdown } from './server/graceful-shutdown.js';
 
 const PORT = process.env.PORT || 8080;
+const SHUTDOWN_TIMEOUT_MS = Number(process.env.SHUTDOWN_TIMEOUT_MS) || 20000;
+
+function withTimeout(
+  label: string,
+  ms: number,
+  run: () => Promise<unknown>
+): Promise<void> {
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  return Promise.race([
+    run(),
+    new Promise<void>((_, reject) => {
+      timer = setTimeout(
+        () => reject(new Error(`${label} timed out after ${ms}ms`)),
+        ms
+      );
+    }),
+  ])
+    .then(() => undefined)
+    .finally(() => {
+      if (timer !== undefined) clearTimeout(timer);
+    });
+}
 
 (async function main() {
   try {
     const context = await buildContext(process.env.PG_CONN);
-    const plugins = await buildPlugins();
+    const { plugins, provider } = await buildPlugins();
     const server = buildServer(context, plugins);
 
     server.listen(PORT, () => {
       console.info(`Server is running on port: ${PORT}`);
     });
 
+    const shutdown = createGracefulShutdown({
+      timeoutMs: SHUTDOWN_TIMEOUT_MS,
+      // Stop accepting connections and wait for in-flight requests to drain.
+      closeServer: () =>
+        new Promise<void>((resolve, reject) => {
+          server.close((error) => (error ? reject(error) : resolve()));
+          // `close` also waits on idle keep-alive sockets, which browser
+          // clients hold open for keepAliveTimeout; dropping them keeps the
+          // drain prompt so the closers below still run inside the timeout.
+          server.closeIdleConnections();
+        }),
+      closers: [
+        // Flush buffered spans without letting an unreachable collector consume
+        // the whole shutdown budget.
+        () =>
+          withTimeout('trace flush', 3000, async () => {
+            if (provider) await provider.shutdown();
+          }),
+        // Close the Postgres pool. postgres.js waits forever by default for
+        // open connections, so bound the wait and let the force timer handle any
+        // still-running queries.
+        () => withTimeout('pg pool close', 5000, () => context.db_client.close()),
+      ],
+    });
+
     ['SIGINT', 'SIGTERM', 'SIGQUIT'].forEach((signal) => {
-      process.on(signal, () => server.close());
+      process.on(signal, () => void shutdown(signal));
     });
 
-    server.on('close', async () => {
-      await context.db_client.close();
-      process.exit(0); // normal termination
+    // Crashes exit non-zero: an exit 0 reads as a clean stop to Kubernetes and
+    // systemd, suppressing restarts and non-zero-exit alerting.
+    process.on('uncaughtException', (error) => {
+      console.error('Uncaught exception:', error);
+      void shutdown('uncaughtException', 1);
+    });
+    process.on('unhandledRejection', (reason) => {
+      console.error('Unhandled rejection:', reason);
+      void shutdown('unhandledRejection', 1);
     });
   } catch (error) {
     console.error('An error occurred:', error);

src/server/graceful-shutdown.ts

@@ -0,0 +1,76 @@
+export { createGracefulShutdown };
+export type { GracefulShutdownOptions };
+
+interface GracefulShutdownOptions {
+  /** Stop accepting connections and resolve once in-flight requests drain. */
+  closeServer: () => Promise<void>;
+  /** Resource teardown to run after the server closes (flush traces, close DB). */
+  closers?: Array<() => Promise<void>>;
+  /** Hard deadline; if draining/teardown exceeds it, force exit. */
+  timeoutMs: number;
+  /** Process exit hook — injectable for tests. */
+  onExit?: (code: number) => void;
+  /** Log sink — injectable for tests. */
+  log?: (message: string, error?: unknown) => void;
+}
+
+/**
+ * Build an idempotent shutdown handler. On the first invocation it drains the
+ * server, runs each closer (a failing closer is logged but doesn't abort the
+ * rest), then exits with the caller's code. A hard timeout guarantees the
+ * process exits even if a connection or teardown step hangs, and `onExit` is
+ * invoked at most once.
+ *
+ * Subsequent invocations (e.g. a second signal) are ignored.
+ */
+function createGracefulShutdown(options: GracefulShutdownOptions) {
+  const {
+    closeServer,
+    closers = [],
+    timeoutMs,
+    onExit = (code) => process.exit(code),
+    log = (message, error) =>
+      error !== undefined ? console.error(message, error) : console.info(message),
+  } = options;
+
+  let started = false;
+
+  /**
+   * `exitCode` is the code used when the drain succeeds; a crash-initiated
+   * shutdown must pass non-zero so supervisors still see a failed exit.
+   */
+  return async function shutdown(reason: string, exitCode = 0): Promise<void> {
+    if (started) return;
+    started = true;
+    log(`Shutting down (${reason})…`);
+
+    let exited = false;
+    const exitOnce = (code: number) => {
+      if (exited) return;
+      exited = true;
+      onExit(code);
+    };
+
+    const forceTimer = setTimeout(() => {
+      log('Graceful shutdown timed out; forcing exit.');
+      exitOnce(1);
+    }, timeoutMs);
+
+    try {
+      await closeServer();
+      for (const close of closers) {
+        try {
+          await close();
+        } catch (error) {
+          log('Error during shutdown step', error);
+        }
+      }
+      exitOnce(exitCode);
+    } catch (error) {
+      log('Error closing server', error);
+      exitOnce(1);
+    } finally {
+      clearTimeout(forceTimer);
+    }
+  };
+}

src/server/plugins.ts

@@ -4,6 +4,7 @@ import { useDisableIntrospection } from '@envelop/disable-introspection';
 import { useOpenTelemetry } from '@envelop/opentelemetry';
 import { inspect } from 'node:util';
 
+import type { BasicTracerProvider } from '@opentelemetry/sdk-trace-base';
 import { initJaegerProvider } from '../tracing/jaeger-tracing.js';
 import { useMetrics } from './metrics.js';
 
@@ -20,8 +21,10 @@ async function buildPlugins() {
 
   plugins.push(useGraphQlJit());
 
+  // Returned so the entry point can flush spans on shutdown.
+  let provider: BasicTracerProvider | undefined;
   if (process.env.ENABLE_LOGGING) {
-    const provider = await initJaegerProvider();
+    provider = await initJaegerProvider();
     plugins.push(
       useOpenTelemetry(
         {
@@ -62,5 +65,5 @@ async function buildPlugins() {
       },
     })
   );
-  return plugins;
+  return { plugins, provider };
 }

tests/unit/graceful-shutdown.test.ts

@@ -0,0 +1,110 @@
+import { describe, test } from 'node:test';
+import assert from 'node:assert';
+import { createGracefulShutdown } from '../../src/server/graceful-shutdown.js';
+
+const silent = () => {};
+
+describe('Graceful shutdown', () => {
+  test('drains the server, runs closers in order, then exits 0', async () => {
+    const calls: string[] = [];
+    const exits: number[] = [];
+    const shutdown = createGracefulShutdown({
+      closeServer: async () => {
+        calls.push('server');
+      },
+      closers: [
+        async () => {
+          calls.push('traces');
+        },
+        async () => {
+          calls.push('db');
+        },
+      ],
+      timeoutMs: 1000,
+      onExit: (code) => exits.push(code),
+      log: silent,
+    });
+
+    await shutdown('SIGTERM');
+    assert.deepStrictEqual(calls, ['server', 'traces', 'db']);
+    assert.deepStrictEqual(exits, [0]);
+  });
+
+  test('is idempotent — a second signal does nothing', async () => {
+    let serverCloses = 0;
+    const exits: number[] = [];
+    const shutdown = createGracefulShutdown({
+      closeServer: async () => {
+        serverCloses += 1;
+      },
+      timeoutMs: 1000,
+      onExit: (code) => exits.push(code),
+      log: silent,
+    });
+
+    await shutdown('SIGTERM');
+    await shutdown('SIGINT');
+    assert.strictEqual(serverCloses, 1);
+    assert.deepStrictEqual(exits, [0]);
+  });
+
+  test('a failing closer is logged but does not abort the rest', async () => {
+    const calls: string[] = [];
+    const exits: number[] = [];
+    const shutdown = createGracefulShutdown({
+      closeServer: async () => {},
+      closers: [
+        async () => {
+          throw new Error('flush failed');
+        },
+        async () => {
+          calls.push('db');
+        },
+      ],
+      timeoutMs: 1000,
+      onExit: (code) => exits.push(code),
+      log: silent,
+    });
+
+    await shutdown('SIGTERM');
+    assert.deepStrictEqual(calls, ['db']);
+    assert.deepStrictEqual(exits, [0]);
+  });
+
+  test('a crash-initiated shutdown drains cleanly but still exits non-zero', async () => {
+    const calls: string[] = [];
+    const exits: number[] = [];
+    const shutdown = createGracefulShutdown({
+      closeServer: async () => {
+        calls.push('server');
+      },
+      closers: [
+        async () => {
+          calls.push('traces');
+        },
+      ],
+      timeoutMs: 1000,
+      onExit: (code) => exits.push(code),
+      log: silent,
+    });
+
+    await shutdown('uncaughtException', 1);
+    // The drain still runs in full — only the exit code differs from a signal.
+    assert.deepStrictEqual(calls, ['server', 'traces']);
+    assert.deepStrictEqual(exits, [1]);
+  });
+
+  test('exits 1 when draining exceeds the timeout, and only once', async () => {
+    const exits: number[] = [];
+    const shutdown = createGracefulShutdown({
+      closeServer: () => new Promise<void>(() => {}), // never resolves
+      timeoutMs: 20,
+      onExit: (code) => exits.push(code),
+      log: silent,
+    });
+
+    shutdown('SIGTERM');
+    await new Promise((resolve) => setTimeout(resolve, 60));
+    assert.deepStrictEqual(exits, [1]);
+  });
+});

tests/unit/metrics.test.ts

@@ -23,10 +23,11 @@ async function graphql(yoga: ReturnType<typeof serverWithFreshMetrics>) {
 }
 
 async function serverWithBuiltPlugins() {
+  const built = await buildPlugins();
   return createYoga({
     schema,
     graphqlEndpoint: '/',
-    plugins: await buildPlugins(),
+    plugins: built.plugins,
   });
 }