@malloy-publisher/server 0.0.234 → 0.0.236

package/src/health.ts

@@ -1,233 +0,0 @@
-/**
- * Health check endpoints and graceful shutdown handling.
- *
- * Provides Kubernetes-compatible health probes:
- * - /health - Overall health status
- * - /health/liveness - Always returns 200 (process is alive)
- * - /health/readiness - Returns 200 when ready, 503 when not ready
- *
- * On SIGTERM, the service enters draining state to allow in-flight requests
- * to complete before shutdown.
- */
-
-import { Express, NextFunction, Request, Response } from "express";
-import { Server } from "http";
-import { components } from "./api";
-import { shutdownSDK } from "./instrumentation";
-import { logger } from "./logger";
-export type OperationalState =
-   components["schemas"]["ServerStatus"]["operationalState"];
-
-let operationalState: OperationalState = "initializing" as OperationalState;
-let ready: boolean = false;
-let preGracefulShutdownCompleted: boolean = false;
-/**
- * Returns the current operational state of the service.
- */
-export function getOperationalState(): OperationalState {
-   return operationalState;
-}
-
-/**
- * Marks the service as ready to serve traffic.
- */
-export function markReady(): void {
-   if (operationalState !== "draining") {
-      operationalState = "serving";
-      ready = true;
-      logger.info("Service marked as ready");
-   } else {
-      logger.error("Service is already draining - cannot mark as ready");
-   }
-}
-
-/**
- * Marks the service as not ready (readiness probe will return 503).
- */
-export function markNotReady(): void {
-   ready = false;
-   logger.info("Service marked as not ready - readiness probe will fail");
-}
-
-/**
- * Registers SIGTERM handler for graceful shutdown.
- *
- * Shutdown sequence:
- * 1. Marks service as not ready (readiness probe returns 503) and enters draining state
- * 2. Waits shutdownDrainDurationSeconds to allow in-flight requests to complete
- * 3. Sets preGracefulShutdownCompleted flag (enables drainingGuard middleware to reject new requests)
- * 4. Closes main server and MCP server (stops accepting new connections)
- * 5. Waits shutdownGracefulCloseTimeoutSeconds (if > 0) for final cleanup
- * 6. Closes logger (last, so any logs emitted during cleanup are flushed)
- * 7. Exits process
- *
- * Note: drainingGuard only rejects requests after step 3 completes. During step 2,
- * the service is draining but still accepts requests (readiness probe returns 503).
- *
- * @param server - Main HTTP server instance
- * @param mcpServer - MCP server instance
- * @param shutdownDrainDurationSeconds - Duration in seconds to wait before closing servers
- * @param shutdownGracefulCloseTimeoutSeconds - Duration in seconds to wait after closing servers before exit
- */
-export function registerSignalHandlers(
-   server: Server,
-   mcpServer: Server,
-   shutdownDrainDurationSeconds: number = 0,
-   shutdownGracefulCloseTimeoutSeconds: number = 0,
-): void {
-   // On SIGTERM, drain in-flight requests before closing the servers (see the
-   // shutdown sequence above). K8s will SIGKILL after terminationGracePeriodSeconds
-   // if this handler does not exit first.
-   process.once("SIGTERM", async () => {
-      logger.info("========== SIGTERM RECEIVED ==========");
-      markNotReady();
-      operationalState = "draining" as OperationalState;
-      logger.info(
-         `Service entering draining state for ${shutdownDrainDurationSeconds} seconds before closing servers...`,
-      );
-
-      await new Promise((resolve) =>
-         setTimeout(() => {
-            preGracefulShutdownCompleted = true;
-            resolve(true);
-         }, shutdownDrainDurationSeconds * 1000),
-      );
-
-      await performGracefulShutdownAfterDrain(
-         server,
-         mcpServer,
-         shutdownGracefulCloseTimeoutSeconds,
-      );
-   });
-}
-
-/**
- * Performs the post-drain shutdown work: closes both HTTP servers,
- * shuts down the OpenTelemetry SDK, waits the optional graceful-close
- * window so any in-flight cleanup can finish logging, closes the
- * winston logger, and exits the process.
- *
- * Exported so unit tests can exercise the close + log + exit ordering
- * without emitting SIGTERM (which would leave module-level
- * operationalState stuck in "draining" and leak into sibling specs).
- */
-export async function performGracefulShutdownAfterDrain(
-   server: Server,
-   mcpServer: Server,
-   shutdownGracefulCloseTimeoutSeconds: number,
-): Promise<void> {
-   const closeServer = (server: Server, name: string) =>
-      new Promise<void>((resolve) => {
-         if (server && server.listening) {
-            server.close((err) => {
-               if (err) {
-                  logger.error(`${name} close error:`, err);
-               } else {
-                  logger.info(`${name} closed`);
-               }
-               resolve();
-            });
-         } else {
-            resolve();
-         }
-      });
-
-   await Promise.all([
-      closeServer(server, "Main server"),
-      closeServer(mcpServer, "MCP server"),
-   ]);
-
-   try {
-      await shutdownSDK();
-      logger.info("OpenTelemetry SDK shut down");
-   } catch (_error) {
-      /* do nothing */
-   }
-
-   try {
-      // Drain in-flight compiles and terminate worker_threads before
-      // we exit so a slow compile doesn't leave orphan worker
-      // processes. Lazy-imported to avoid pulling the pool module
-      // into the health.ts dep graph for tests that don't exercise
-      // the compile path.
-      const { getPackageLoadPool } = await import(
-         "./package_load/package_load_pool"
-      );
-      await getPackageLoadPool().shutdown();
-      logger.info("Package-load worker pool shut down");
-   } catch (_error) {
-      /* do nothing */
-   }
-
-   if (shutdownGracefulCloseTimeoutSeconds > 0) {
-      logger.info(
-         `Waiting ${shutdownGracefulCloseTimeoutSeconds} seconds after server close before exit...`,
-      );
-      await new Promise((resolve) =>
-         setTimeout(resolve, shutdownGracefulCloseTimeoutSeconds * 1000),
-      );
-   }
-
-   // Close the logger last so anything emitted during the wait window
-   // above (or by other shutdown paths still running) reaches its
-   // transports. Closing earlier triggers winston's
-   // "Attempt to write logs with no transports" warning on any
-   // subsequent logger call.
-   try {
-      logger.close();
-   } catch (_error) {
-      /* do nothing */
-   }
-
-   process.exit(0);
-}
-/**
- * Middleware that returns 503 for non-health and metrics requests when service is draining.
- * Must be registered before application routes.
- */
-export function drainingGuard(
-   req: Request,
-   res: Response,
-   next: NextFunction,
-): void {
-   if (
-      operationalState === "draining" &&
-      preGracefulShutdownCompleted &&
-      !req.path.startsWith("/health") &&
-      !req.path.startsWith("/metrics")
-   ) {
-      res.status(503).json({
-         message: "Service is draining",
-         details: "The service is shutting down and hence is not available.",
-      });
-      return;
-   }
-   next();
-}
-
-/**
- * Registers health check endpoints
- */
-export function registerHealthEndpoints(app: Express): void {
-   app.get("/health", (_req: Request, res: Response) => {
-      res.status(200).json({
-         status: "UP",
-         components: {
-            livenessState: { status: "UP" },
-            readinessState: { status: ready ? "UP" : "DOWN" },
-         },
-         groups: ["livenessState", "readinessState"],
-      });
-   });
-
-   app.get("/health/liveness", (_req: Request, res: Response) => {
-      res.status(200).json({ status: "UP" });
-   });
-
-   app.get("/health/readiness", (_req: Request, res: Response) => {
-      const isReady = ready;
-      res.status(isReady ? 200 : 503).json({
-         status: isReady ? "UP" : "DOWN",
-      });
-   });
-}

package/src/heap_check.spec.ts

@@ -1,144 +0,0 @@
-import { afterEach, beforeEach, describe, expect, it, mock } from "bun:test";
-
-import {
-   checkHeapConfiguration,
-   resetHeapTelemetryForTesting,
-} from "./heap_check";
-import {
-   startMetricsHarness,
-   type MetricsHarness,
-} from "./test_helpers/metrics_harness";
-
-function makeLogStub(): {
-   warn: ReturnType<typeof mock>;
-   info: ReturnType<typeof mock>;
-} {
-   return {
-      warn: mock(() => undefined),
-      info: mock(() => undefined),
-   };
-}
-
-describe("checkHeapConfiguration", () => {
-   it("warns + reports warned=true when heap limit is below the 2 GiB threshold", () => {
-      const log = makeLogStub();
-      // 1.5 GiB — below the recommended floor.
-      const { warned } = checkHeapConfiguration({
-         getHeapStatistics: () => ({
-            heap_size_limit: 1.5 * 1024 * 1024 * 1024,
-         }),
-         log,
-      });
-      expect(warned).toBe(true);
-      expect(log.warn).toHaveBeenCalledTimes(1);
-      expect(log.info).not.toHaveBeenCalled();
-      const args = log.warn.mock.calls[0] as unknown as [
-         string,
-         Record<string, unknown>,
-      ];
-      // Operator must be able to grep for the offending env-var
-      // hint without guessing the surrounding sentence.
-      expect(args[0]).toContain("--max-old-space-size");
-      expect(args[0]).toContain("MiB");
-      expect(args[1].heapSizeLimitBytes).toBe(1.5 * 1024 * 1024 * 1024);
-   });
-
-   it("logs at info + reports warned=false when heap limit meets the recommendation", () => {
-      const log = makeLogStub();
-      const { warned, heapSizeLimitBytes } = checkHeapConfiguration({
-         getHeapStatistics: () => ({
-            heap_size_limit: 4 * 1024 * 1024 * 1024,
-         }),
-         log,
-      });
-      expect(warned).toBe(false);
-      expect(heapSizeLimitBytes).toBe(4 * 1024 * 1024 * 1024);
-      expect(log.warn).not.toHaveBeenCalled();
-      expect(log.info).toHaveBeenCalledTimes(1);
-   });
-
-   it("treats exactly-the-recommended-value as 'meets threshold' (no warn)", () => {
-      // Boundary: the comparison is `<`, so equality should pass.
-      // Lock the boundary so a future tightening to `<=` requires
-      // an explicit test change rather than silently regressing.
-      const log = makeLogStub();
-      const { warned } = checkHeapConfiguration({
-         getHeapStatistics: () => ({
-            heap_size_limit: 2 * 1024 * 1024 * 1024,
-         }),
-         log,
-      });
-      expect(warned).toBe(false);
-      expect(log.warn).not.toHaveBeenCalled();
-   });
-
-   it("does not throw or warn under tiny heaps; the cap helpers still work", () => {
-      // Pathological "this pod has 128 MiB" case: we want a noisy
-      // warning, not a process crash, so the publisher still boots
-      // and the operator sees the message in pod logs.
-      const log = makeLogStub();
-      expect(() =>
-         checkHeapConfiguration({
-            getHeapStatistics: () => ({ heap_size_limit: 128 * 1024 * 1024 }),
-            log,
-         }),
-      ).not.toThrow();
-      expect(log.warn).toHaveBeenCalledTimes(1);
-   });
-
-   it("uses live v8.getHeapStatistics when no override is provided (smoke check)", () => {
-      // No assertion on warn/info — the result depends on how Node
-      // was started — just that the call resolves and returns a
-      // sensible structure. Locks the production code path against
-      // accidental coupling to the injected getter.
-      const result = checkHeapConfiguration();
-      expect(result.heapSizeLimitBytes).toBeGreaterThan(0);
-      expect(typeof result.warned).toBe("boolean");
-   });
-
-   describe("telemetry", () => {
-      let harness: MetricsHarness;
-
-      beforeEach(async () => {
-         harness = await startMetricsHarness();
-         resetHeapTelemetryForTesting();
-      });
-
-      afterEach(async () => {
-         resetHeapTelemetryForTesting();
-         await harness.shutdown();
-      });
-
-      it("publisher_heap_size_limit_bytes reports the live V8 heap_size_limit", async () => {
-         const log = makeLogStub();
-         checkHeapConfiguration({ log });
-         const value = await harness.collectGauge(
-            "publisher_heap_size_limit_bytes",
-         );
-         // Live value — we just assert it's a sensible positive
-         // number so the test isn't sensitive to how this Bun
-         // process was launched.
-         expect(typeof value).toBe("number");
-         expect(value).toBeGreaterThan(0);
-      });
-
-      it("publisher_heap_used_bytes reports the live V8 used_heap_size", async () => {
-         const log = makeLogStub();
-         checkHeapConfiguration({ log });
-         const value = await harness.collectGauge("publisher_heap_used_bytes");
-         expect(typeof value).toBe("number");
-         expect(value).toBeGreaterThan(0);
-      });
-
-      it("does not register the gauges before checkHeapConfiguration is called (lazy install)", async () => {
-         // The gauges are wired up via `installHeapGauges`, which
-         // is intentionally called from `checkHeapConfiguration`
-         // so the OTel SDK is fully up before instruments are
-         // registered. Without that, instruments would bind to
-         // NoOp during module load and never emit data.
-         expect(
-            await harness.collectGauge("publisher_heap_size_limit_bytes"),
-         ).toBeUndefined();
-      });
-   });
-});

package/src/heap_check.ts

@@ -1,144 +0,0 @@
-import { publisherMeter } from "./telemetry";
-import * as v8 from "v8";
-
-import { logger } from "./logger";
-
-/**
- * Subset of the winston logger surface this module needs. Defined
- * structurally so tests can pass a plain object stub and so we
- * don't take a direct dependency on the concrete winston type.
- */
-interface HeapCheckLogger {
-   warn: (...args: unknown[]) => unknown;
-   info: (...args: unknown[]) => unknown;
-}
-
-/**
- * Observable gauges for heap configuration so dashboards can render
- * "configured heap ceiling" alongside the RSS / back-pressure
- * timeseries from `PackageMemoryGovernor`. The values are observed
- * on every scrape rather than cached — `v8.getHeapStatistics()` is
- * cheap (a single VM call) and serving the live value avoids stale
- * reads.
- *
- * Lazy init for the same reason as `query_timeout.ts`: instruments
- * captured before `setGlobalMeterProvider` are bound to NoOp. We
- * call `installHeapGauges` from `checkHeapConfiguration` so it
- * runs once at startup, after the OTel SDK is up.
- *
- * Mirrors the governor's `publisher_*` unlabeled style.
- */
-let heapGaugesInstalled = false;
-function installHeapGauges(): void {
-   if (heapGaugesInstalled) return;
-   heapGaugesInstalled = true;
-   const meter = publisherMeter();
-   meter
-      .createObservableGauge("publisher_heap_size_limit_bytes", {
-         description:
-            "V8 heap_size_limit (--max-old-space-size). Compare with PUBLISHER_MAX_MEMORY_BYTES.",
-         unit: "By",
-      })
-      .addCallback((observation) => {
-         observation.observe(v8.getHeapStatistics().heap_size_limit);
-      });
-   meter
-      .createObservableGauge("publisher_heap_used_bytes", {
-         description:
-            "Current V8 used_heap_size in bytes. Watch this alongside publisher_process_rss_bytes; the two diverge under native-allocator pressure (DuckDB, etc.).",
-         unit: "By",
-      })
-      .addCallback((observation) => {
-         observation.observe(v8.getHeapStatistics().used_heap_size);
-      });
-}
-
-/**
- * Visible for tests; production code never calls this. Resets the
- * lazy guard so a re-installation captures into a fresh
- * MeterProvider.
- */
-export function resetHeapTelemetryForTesting(): void {
-   heapGaugesInstalled = false;
-}
-
-/**
- * Minimum V8 heap ceiling (`--max-old-space-size`) the publisher
- * expects in production. Below this the row/byte caps from earlier
- * steps still apply, but a single buffered model query at the
- * default 50 MB byte cap plus the surrounding `Result` allocation
- * can plausibly chew through the remaining headroom and OOM the
- * process before back-pressure trips.
- *
- * 2 GiB is the smallest value at which the defaults are comfortably
- * survivable. Operators running explicitly tuned-down pods (e.g. a
- * lightweight smoke-test deploy) can ignore the warning.
- */
-const MIN_RECOMMENDED_HEAP_BYTES = 2 * 1024 * 1024 * 1024;
-
-/**
- * Probe `v8.getHeapStatistics().heap_size_limit` at startup and
- * emit a single structured warning when the process is configured
- * with a heap ceiling below {@link MIN_RECOMMENDED_HEAP_BYTES}.
- *
- * Why warn (not exit):
- * - Smaller pods are legitimate (CI, local dev, smoke tests).
- *   Hard-failing them would be hostile to those workflows.
- * - The earlier steps already bound memory growth per request; this
- *   is a "you probably want to know" signal, not a safety
- *   interlock.
- * - A warning at boot is grep-able in pod logs / dashboards and
- *   surfaces faster than waiting for the first OOMKill.
- *
- * Returns the observed heap limit so the caller (server.ts startup)
- * can also surface it in startup metrics if desired.
- */
-/**
- * Test seam: parameters allow injecting a fake heap-stats getter
- * and logger so the bun/sinon "ES modules cannot be stubbed"
- * restriction doesn't force a module-level mock. Production calls
- * (server.ts startup) use the defaults; tests pass in stubs.
- */
-export interface CheckHeapOptions {
-   getHeapStatistics?: () => Pick<v8.HeapInfo, "heap_size_limit">;
-   log?: HeapCheckLogger;
-}
-
-export function checkHeapConfiguration(options: CheckHeapOptions = {}): {
-   heapSizeLimitBytes: number;
-   warned: boolean;
-} {
-   // Install heap-related observable gauges on the same startup
-   // tick. Idempotent — re-calling in tests / warmup paths is
-   // safe and re-uses the same instruments.
-   installHeapGauges();
-   const getStats = options.getHeapStatistics ?? v8.getHeapStatistics;
-   const log = options.log ?? logger;
-   const stats = getStats();
-   const heapSizeLimitBytes = stats.heap_size_limit;
-   const limitMiB = Math.round(heapSizeLimitBytes / (1024 * 1024));
-   const recommendedMiB = Math.round(
-      MIN_RECOMMENDED_HEAP_BYTES / (1024 * 1024),
-   );
-
-   if (heapSizeLimitBytes < MIN_RECOMMENDED_HEAP_BYTES) {
-      log.warn(
-         `V8 heap_size_limit is ${limitMiB} MiB, below the recommended ${recommendedMiB} MiB. ` +
-            `Pass --max-old-space-size=${recommendedMiB} (or higher) on the node process to keep ` +
-            `the row/byte caps (PUBLISHER_MAX_QUERY_ROWS / PUBLISHER_MAX_RESPONSE_BYTES) within ` +
-            `safe margin. With a smaller heap, a single large query can OOM the pod before the ` +
-            `memory governor's back-pressure has a chance to trip.`,
-         {
-            heapSizeLimitBytes,
-            recommendedHeapSizeBytes: MIN_RECOMMENDED_HEAP_BYTES,
-         },
-      );
-      return { heapSizeLimitBytes, warned: true };
-   }
-
-   log.info(
-      `V8 heap_size_limit is ${limitMiB} MiB (>= recommended ${recommendedMiB} MiB).`,
-      { heapSizeLimitBytes },
-   );
-   return { heapSizeLimitBytes, warned: false };
-}