@malloy-publisher/server 0.0.198 → 0.0.200

package/src/health.ts

@@ -57,8 +57,8 @@ export function markNotReady(): void {
  * 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. Closes logger
- * 6. Waits shutdownGracefulCloseTimeoutSeconds (if > 0) for final cleanup
+ * 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,
@@ -92,51 +92,94 @@ export function registerSignalHandlers(
          }, shutdownDrainDurationSeconds * 1000),
       );
 
-      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 {
-         logger.close();
-      } 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),
-         );
-      }
-      process.exit(0);
+      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.

package/src/heap_check.spec.ts

@@ -0,0 +1,144 @@
+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

@@ -0,0 +1,144 @@
+import { metrics } from "@opentelemetry/api";
+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 = metrics.getMeter("publisher");
+   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 };
+}

package/src/instrumentation.ts

@@ -1,3 +1,4 @@
+import { monitorEventLoopDelay } from "node:perf_hooks";
 import { metrics } from "@opentelemetry/api";
 import { getNodeAutoInstrumentations } from "@opentelemetry/auto-instrumentations-node";
 import { OTLPLogExporter } from "@opentelemetry/exporter-logs-otlp-proto";
@@ -116,6 +117,55 @@ const httpRequestCount = meter.createCounter("http_server_requests_total", {
    description: "Total number of HTTP requests",
 });
 
+// Event-loop-delay metrics. A blocked event loop is the only way the
+// /health/liveness probe (a pure synchronous 200 handler) can fail under K8s,
+// so we surface p50/p99/max so an operator can correlate liveness restarts
+// with sustained event-loop pressure (large Malloy compiles, GC, etc.).
+const eventLoopHistogram = monitorEventLoopDelay({ resolution: 20 });
+eventLoopHistogram.enable();
+
+const eventLoopLagP50 = meter.createObservableGauge(
+   "publisher_event_loop_lag_p50_ms",
+   {
+      description:
+         "Event loop delay p50 since the last scrape, in milliseconds",
+      unit: "ms",
+   },
+);
+const eventLoopLagP99 = meter.createObservableGauge(
+   "publisher_event_loop_lag_p99_ms",
+   {
+      description:
+         "Event loop delay p99 since the last scrape, in milliseconds",
+      unit: "ms",
+   },
+);
+const eventLoopLagMax = meter.createObservableGauge(
+   "publisher_event_loop_lag_max_ms",
+   {
+      description:
+         "Event loop delay max since the last scrape, in milliseconds",
+      unit: "ms",
+   },
+);
+
+// Sample all three in one batch so the histogram reset can't race the reads.
+meter.addBatchObservableCallback(
+   (observableResult) => {
+      observableResult.observe(
+         eventLoopLagP50,
+         eventLoopHistogram.percentile(50) / 1e6,
+      );
+      observableResult.observe(
+         eventLoopLagP99,
+         eventLoopHistogram.percentile(99) / 1e6,
+      );
+      observableResult.observe(eventLoopLagMax, eventLoopHistogram.max / 1e6);
+      eventLoopHistogram.reset();
+   },
+   [eventLoopLagP50, eventLoopLagP99, eventLoopLagMax],
+);
+
 const IGNORED_PATHS = new Set([
    "/health",
    "/health/liveness",

package/src/mcp/handler_utils.ts

@@ -7,6 +7,7 @@ import {
    ModelNotFoundError,
    ModelCompilationError,
    EnvironmentNotFoundError,
+   ServiceUnavailableError,
 } from "../errors";
 import {
    getNotFoundError,
@@ -132,6 +133,9 @@ export async function getModelForQuery(
          environmentName,
          false,
       );
+      // Shed load before any disk / DB work; mirrors the HTTP query
+      // controllers so MCP traffic obeys the same back-pressure rules.
+      environment.assertCanAdmitQuery();
       const pkg = await environment.getPackage(packageName, false);
       const model = pkg.getModel(modelPath);
       if (!model || model.getModelType() === "notebook") {
@@ -163,6 +167,16 @@ export async function getModelForQuery(
             `${environmentName}/${packageName}/${modelPath}`,
             error,
          );
+      } else if (error instanceof ServiceUnavailableError) {
+         // Back-pressure: don't dress this up as a 404/500. Surface the
+         // server's own message so the MCP caller knows to retry.
+         errorDetails = {
+            message: error.message,
+            suggestions: [
+               "Retry after the publisher's memory usage drops below the configured low-water mark.",
+               "If this happens repeatedly, raise PUBLISHER_MAX_MEMORY_BYTES or scale up the pod.",
+            ],
+         } satisfies ErrorDetails;
       } else {
          // Unexpected error during setup
          errorDetails = getInternalError("executeQuery (Setup)", error);

package/src/mcp/tools/execute_query_tool.ts

@@ -1,7 +1,14 @@
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
 import { z } from "zod";
+import type { GivenValue } from "@malloydata/malloy";
+import { getQueryTimeoutMs } from "../../config";
 import { logger } from "../../logger";
+import {
+   tryAcquireQuerySlot,
+   type QuerySlotHandle,
+} from "../../query_concurrency";
+import { runWithQueryTimeout } from "../../query_timeout";
 import { EnvironmentStore } from "../../service/environment_store";
 import { getMalloyErrorDetails, type ErrorDetails } from "../error_messages";
 import { buildMalloyUri, getModelForQuery } from "../handler_utils";
@@ -30,6 +37,12 @@ const executeQueryShape = {
       .describe(
          "Filter parameter values keyed by filter name. Used with sources that declare #(filter) annotations.",
       ),
+   givens: z
+      .record(z.unknown())
+      .optional()
+      .describe(
+         "Per-query given values that override model defaults. Keys are given names declared in the model's given: block.",
+      ),
 };
 
 // Type inference is handled automatically by the MCP server based on the executeQueryShape
@@ -56,6 +69,7 @@ export function registerExecuteQueryTool(
             sourceName,
             queryName,
             filterParams,
+            givens,
          } = params;
 
          logger.info("[MCP Tool executeQuery] Received params:", { params });
@@ -120,14 +134,30 @@ export function registerExecuteQueryTool(
          logger.info(
             `[MCP Tool executeQuery] Model found. Proceeding to execute query.`,
          );
+         // Per-pod concurrency slot. MCP shares the same slot pool
+         // as the HTTP query routes so a hot agent loop can't
+         // bypass PUBLISHER_MAX_CONCURRENT_QUERIES. `mcp:executeQuery`
+         // is a fixed label so the dashboard can separate MCP load
+         // from HTTP route load. Acquisition can throw
+         // ServiceUnavailableError; the existing catch below surfaces
+         // it as the standard MCP error-content payload.
+         let querySlot: QuerySlotHandle | null = null;
          try {
+            querySlot = tryAcquireQuerySlot("mcp:executeQuery");
             // If ad-hoc query is provided, use it directly in the 3rd arg
             if (query) {
-               const { result } = await model.getQueryResults(
-                  undefined,
-                  undefined,
-                  query,
-                  filterParams,
+               const { result } = await runWithQueryTimeout(
+                  (abortSignal) =>
+                     model.getQueryResults(
+                        undefined,
+                        undefined,
+                        query,
+                        filterParams,
+                        undefined,
+                        givens as Record<string, GivenValue> | undefined,
+                        abortSignal,
+                     ),
+                  getQueryTimeoutMs(),
                );
                const { validateRenderTags } = await import(
                   "@malloydata/render-validator"
@@ -169,11 +199,18 @@ export function registerExecuteQueryTool(
 
                return { isError: false, content };
             } else if (queryName) {
-               const { result } = await model.getQueryResults(
-                  sourceName,
-                  queryName,
-                  undefined,
-                  filterParams,
+               const { result } = await runWithQueryTimeout(
+                  (abortSignal) =>
+                     model.getQueryResults(
+                        sourceName,
+                        queryName,
+                        undefined,
+                        filterParams,
+                        undefined,
+                        givens as Record<string, GivenValue> | undefined,
+                        abortSignal,
+                     ),
+                  getQueryTimeoutMs(),
                );
                const { validateRenderTags } = await import(
                   "@malloydata/render-validator"
@@ -259,6 +296,11 @@ export function registerExecuteQueryTool(
                   },
                ],
             };
+         } finally {
+            // Release on every exit path — success, error, or
+            // unreachable code-path throw. `release()` is idempotent
+            // so a double-fault during cleanup can't double-decrement.
+            querySlot?.release();
          }
       },
    );