@malloy-publisher/server 0.0.198-dev → 0.0.198-dev2

package/src/service/package_memory_governor.spec.ts

@@ -0,0 +1,173 @@
+import { describe, expect, it } from "bun:test";
+
+import type { MemoryGovernorConfig } from "../config";
+import { PackageMemoryGovernor } from "./package_memory_governor";
+
+const ONE_GB = 1024 * 1024 * 1024;
+
+function makeConfig(
+   overrides: Partial<MemoryGovernorConfig> = {},
+): MemoryGovernorConfig {
+   return {
+      maxMemoryBytes: ONE_GB,
+      highWaterFraction: 0.8,
+      lowWaterFraction: 0.7,
+      checkIntervalMs: 5_000,
+      backpressureEnabled: true,
+      ...overrides,
+   };
+}
+
+/**
+ * Test driver that lets us push a sequence of RSS values into the
+ * governor and inspect the state machine's reactions deterministically
+ * — no real allocations, no real timers.
+ */
+class FakeRssSampler {
+   private value = 0;
+   constructor(initial = 0) {
+      this.value = initial;
+   }
+   set(value: number): void {
+      this.value = value;
+   }
+   sampler = (): number => this.value;
+}
+
+describe("PackageMemoryGovernor", () => {
+   it("does not activate back-pressure below the high-water mark", () => {
+      const rss = new FakeRssSampler(0.5 * ONE_GB);
+      const gov = new PackageMemoryGovernor(makeConfig(), rss.sampler);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(false);
+   });
+
+   it("activates back-pressure at or above the high-water mark", () => {
+      const rss = new FakeRssSampler(0);
+      const gov = new PackageMemoryGovernor(makeConfig(), rss.sampler);
+
+      rss.set(0.79 * ONE_GB);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(false);
+
+      // 0.8 * 1GB is exactly the high-water threshold; using >= so it
+      // trips on the boundary.
+      rss.set(0.8 * ONE_GB);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(true);
+   });
+
+   it("does not clear back-pressure inside the hysteresis band", () => {
+      const rss = new FakeRssSampler(0.9 * ONE_GB);
+      const gov = new PackageMemoryGovernor(makeConfig(), rss.sampler);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(true);
+
+      // Between low (0.7) and high (0.8) — must stay backpressured.
+      rss.set(0.75 * ONE_GB);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(true);
+   });
+
+   it("clears back-pressure at or below the low-water mark", () => {
+      const rss = new FakeRssSampler(0.9 * ONE_GB);
+      const gov = new PackageMemoryGovernor(makeConfig(), rss.sampler);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(true);
+
+      // The implementation floors lowWaterBytes (= 0.7 * 1GB → 751619276),
+      // so we need to feed a value at or below that integer — `0.7 * 1GB`
+      // as a float is 751619276.8 which sits just above the threshold.
+      rss.set(0.69 * ONE_GB);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(false);
+   });
+
+   it("re-activates after recovery if RSS climbs again", () => {
+      const rss = new FakeRssSampler(0);
+      const gov = new PackageMemoryGovernor(makeConfig(), rss.sampler);
+
+      rss.set(0.85 * ONE_GB);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(true);
+
+      rss.set(0.6 * ONE_GB);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(false);
+
+      rss.set(0.9 * ONE_GB);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(true);
+   });
+
+   it("samples but never flips the flag when backpressureEnabled=false", () => {
+      const rss = new FakeRssSampler(0.95 * ONE_GB);
+      const gov = new PackageMemoryGovernor(
+         makeConfig({ backpressureEnabled: false }),
+         rss.sampler,
+      );
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(false);
+      // Status still tracks RSS even though the flag is suppressed.
+      expect(gov.getStatus().rssBytes).toBe(0.95 * ONE_GB);
+   });
+
+   it("survives a throwing sampler without crashing or flipping state", () => {
+      let throwOnce = true;
+      const gov = new PackageMemoryGovernor(makeConfig(), () => {
+         if (throwOnce) {
+            throwOnce = false;
+            throw new Error("simulated sampling failure");
+         }
+         return 0.4 * ONE_GB;
+      });
+
+      // First tick: sampler throws; governor swallows it and leaves
+      // the state untouched.
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(false);
+
+      // Second tick succeeds.
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(false);
+   });
+
+   it("start() takes an immediate sample so a hot-start respects the cap", () => {
+      const rss = new FakeRssSampler(0.95 * ONE_GB);
+      const gov = new PackageMemoryGovernor(
+         // Big interval so we know the initial sample isn't from a
+         // delayed tick.
+         makeConfig({ checkIntervalMs: 60_000 }),
+         rss.sampler,
+      );
+      gov.start();
+      expect(gov.isBackpressured()).toBe(true);
+      gov.stop();
+   });
+
+   it("stop() clears back-pressure and is idempotent", () => {
+      const rss = new FakeRssSampler(0.95 * ONE_GB);
+      const gov = new PackageMemoryGovernor(makeConfig(), rss.sampler);
+      gov.tick();
+      expect(gov.isBackpressured()).toBe(true);
+
+      gov.stop();
+      expect(gov.isBackpressured()).toBe(false);
+      // Second call is a no-op (no thrown error, flag stays cleared).
+      gov.stop();
+      expect(gov.isBackpressured()).toBe(false);
+   });
+
+   it("exposes computed threshold bytes through getStatus", () => {
+      const rss = new FakeRssSampler(0.4 * ONE_GB);
+      const gov = new PackageMemoryGovernor(makeConfig(), rss.sampler);
+      gov.tick();
+      const status = gov.getStatus();
+      expect(status.maxMemoryBytes).toBe(ONE_GB);
+      expect(status.highWaterBytes).toBe(Math.floor(0.8 * ONE_GB));
+      expect(status.lowWaterBytes).toBe(Math.floor(0.7 * ONE_GB));
+      expect(status.rssBytes).toBe(0.4 * ONE_GB);
+      expect(status.backpressured).toBe(false);
+      expect(typeof status.lastSampledAt).toBe("number");
+   });
+});

package/src/service/package_memory_governor.ts

@@ -0,0 +1,233 @@
+import { metrics } from "@opentelemetry/api";
+
+import type { MemoryGovernorConfig } from "../config";
+import { logger } from "../logger";
+
+/**
+ * Snapshot returned by {@link PackageMemoryGovernor.getStatus} for
+ * health endpoints, tests, and ad-hoc logging.
+ */
+export interface MemoryGovernorStatus {
+   rssBytes: number;
+   maxMemoryBytes: number;
+   highWaterBytes: number;
+   lowWaterBytes: number;
+   backpressured: boolean;
+   /** Wall-clock ms of the last successful RSS sample. */
+   lastSampledAt: number | null;
+}
+
+/**
+ * Function that returns the current process RSS in bytes. Injectable
+ * so unit tests can drive the governor with a deterministic source
+ * without spinning real allocations.
+ */
+export type RssSampler = () => number;
+
+const DEFAULT_RSS_SAMPLER: RssSampler = () => process.memoryUsage().rss;
+
+/**
+ * Polls process RSS on a fixed interval and toggles a single
+ * `backpressured` flag using a low/high-water hysteresis band:
+ *
+ *  - RSS >= highWater → set `backpressured = true`
+ *  - RSS <= lowWater  → set `backpressured = false`
+ *  - in between       → leave the flag unchanged
+ *
+ * Controllers consult {@link isBackpressured} on hot paths that would
+ * load a *new* package into memory (`addPackage`, reload, install) and
+ * throw `ServiceUnavailableError` so the request fails fast as 503
+ * instead of pushing the pod into an OOM kill.
+ *
+ * Already-loaded packages remain fully serviceable while back-pressure
+ * is active — this is admission control on new memory, not a cache
+ * eviction. Recovery happens naturally as in-flight traffic completes
+ * and the kernel reclaims pages.
+ *
+ * Disabled by default; only constructed when
+ * `getMemoryGovernorConfig()` returns a non-null config (driven by
+ * `PUBLISHER_MAX_MEMORY_BYTES`).
+ */
+export class PackageMemoryGovernor {
+   private readonly config: MemoryGovernorConfig;
+   private readonly rssSampler: RssSampler;
+   private readonly highWaterBytes: number;
+   private readonly lowWaterBytes: number;
+   private timer: ReturnType<typeof setInterval> | null = null;
+   private backpressured = false;
+   private lastSampledRss = 0;
+   private lastSampledAt: number | null = null;
+   private readonly backpressureActivationsCounter: ReturnType<
+      ReturnType<typeof metrics.getMeter>["createCounter"]
+   >;
+
+   constructor(config: MemoryGovernorConfig, rssSampler?: RssSampler) {
+      this.config = config;
+      this.rssSampler = rssSampler ?? DEFAULT_RSS_SAMPLER;
+      this.highWaterBytes = Math.floor(
+         config.maxMemoryBytes * config.highWaterFraction,
+      );
+      this.lowWaterBytes = Math.floor(
+         config.maxMemoryBytes * config.lowWaterFraction,
+      );
+
+      const meter = metrics.getMeter("publisher");
+
+      // Periodic gauge: current process RSS in bytes.
+      meter
+         .createObservableGauge("publisher_process_rss_bytes", {
+            description:
+               "Current resident set size of the publisher process in bytes",
+            unit: "By",
+         })
+         .addCallback((observation) => {
+            observation.observe(this.rssSampler());
+         });
+
+      // Periodic gauge: 1 when admission control is rejecting new
+      // package loads, 0 otherwise.
+      meter
+         .createObservableGauge("publisher_memory_backpressure_active", {
+            description:
+               "1 when the publisher is rejecting new package loads to stay under PUBLISHER_MAX_MEMORY_BYTES; 0 otherwise",
+         })
+         .addCallback((observation) => {
+            observation.observe(this.backpressured ? 1 : 0);
+         });
+
+      // Cumulative counter for how many times we have transitioned
+      // from `false → true`. Useful for alerting on a flapping pod.
+      this.backpressureActivationsCounter = meter.createCounter(
+         "publisher_memory_backpressure_activations_total",
+         {
+            description:
+               "Number of times the memory governor has activated back-pressure",
+         },
+      );
+
+      // Static gauges so dashboards can render the band alongside RSS
+      // without needing to plumb config separately.
+      meter
+         .createObservableGauge("publisher_memory_max_bytes", {
+            description: "Configured PUBLISHER_MAX_MEMORY_BYTES",
+            unit: "By",
+         })
+         .addCallback((observation) =>
+            observation.observe(this.config.maxMemoryBytes),
+         );
+      meter
+         .createObservableGauge("publisher_memory_high_water_bytes", {
+            description: "RSS threshold at which back-pressure activates",
+            unit: "By",
+         })
+         .addCallback((observation) =>
+            observation.observe(this.highWaterBytes),
+         );
+      meter
+         .createObservableGauge("publisher_memory_low_water_bytes", {
+            description: "RSS threshold at which back-pressure clears",
+            unit: "By",
+         })
+         .addCallback((observation) => observation.observe(this.lowWaterBytes));
+   }
+
+   /**
+    * Begin periodic RSS sampling. Safe to call multiple times — extra
+    * calls are no-ops. The interval is `.unref()`'d so the governor
+    * does not keep the process alive on its own.
+    */
+   public start(): void {
+      if (this.timer !== null) return;
+      // Take an immediate sample so a freshly-started server with
+      // pre-existing high RSS goes into back-pressure right away
+      // instead of waiting `checkIntervalMs` for the first tick.
+      this.tick();
+      this.timer = setInterval(() => this.tick(), this.config.checkIntervalMs);
+      // Tolerate environments without Timer#unref (e.g. some bundlers).
+      (
+         this.timer as ReturnType<typeof setInterval> & {
+            unref?: () => void;
+         }
+      ).unref?.();
+      logger.info(
+         `PackageMemoryGovernor started (max=${this.config.maxMemoryBytes}B, high=${this.highWaterBytes}B, low=${this.lowWaterBytes}B, interval=${this.config.checkIntervalMs}ms, backpressure=${this.config.backpressureEnabled})`,
+      );
+   }
+
+   /**
+    * Stop the periodic sampler. Idempotent. Clears the back-pressure
+    * flag so any in-process logic that consults
+    * {@link isBackpressured} during shutdown sees a permissive state.
+    */
+   public stop(): void {
+      if (this.timer !== null) {
+         clearInterval(this.timer);
+         this.timer = null;
+      }
+      this.backpressured = false;
+   }
+
+   /**
+    * Sample RSS once and apply the hysteresis band. Exposed (rather
+    * than kept private) so callers can force a fresh check right
+    * after they finish loading a new package, and so tests can drive
+    * the governor synchronously.
+    */
+   public tick(): void {
+      let rss: number;
+      try {
+         rss = this.rssSampler();
+      } catch (err) {
+         // Sampling failures must never crash the server. Log and
+         // skip; the next interval will retry. Leave the flag
+         // unchanged so we neither over- nor under-react to a single
+         // measurement glitch.
+         logger.error("PackageMemoryGovernor: RSS sample failed", {
+            error: err,
+         });
+         return;
+      }
+      this.lastSampledRss = rss;
+      this.lastSampledAt = Date.now();
+
+      if (!this.config.backpressureEnabled) {
+         // Feature dial: keep sampling for metrics but never flip
+         // the flag. Useful for monitoring-only rollouts before
+         // enabling the actual 503 behaviour.
+         return;
+      }
+
+      if (rss >= this.highWaterBytes && !this.backpressured) {
+         this.backpressured = true;
+         this.backpressureActivationsCounter.add(1);
+         logger.warn(
+            `PackageMemoryGovernor: activating back-pressure (rss=${rss}B >= high=${this.highWaterBytes}B). New package loads will be rejected with HTTP 503 until rss <= ${this.lowWaterBytes}B.`,
+         );
+      } else if (rss <= this.lowWaterBytes && this.backpressured) {
+         this.backpressured = false;
+         logger.info(
+            `PackageMemoryGovernor: clearing back-pressure (rss=${rss}B <= low=${this.lowWaterBytes}B).`,
+         );
+      }
+   }
+
+   /**
+    * True iff new package-load requests should be rejected with HTTP
+    * 503. Cheap O(1) read of a private boolean; safe to call on every
+    * request.
+    */
+   public isBackpressured(): boolean {
+      return this.backpressured;
+   }
+
+   public getStatus(): MemoryGovernorStatus {
+      return {
+         rssBytes: this.lastSampledRss,
+         maxMemoryBytes: this.config.maxMemoryBytes,
+         highWaterBytes: this.highWaterBytes,
+         lowWaterBytes: this.lowWaterBytes,
+         backpressured: this.backpressured,
+         lastSampledAt: this.lastSampledAt,
+      };
+   }
+}

package/src/service/package_race.spec.ts

@@ -0,0 +1,208 @@
+import { afterEach, beforeEach, describe, expect, it } from "bun:test";
+import * as fs from "fs/promises";
+import * as os from "os";
+import * as path from "path";
+import { Environment } from "./environment";
+
+/**
+ * Race-condition regression tests for the package-directory pipeline.
+ *
+ * Three tests, all deterministic without timing-based flakiness:
+ *
+ *  1. **Behavioral race repro** — concurrently install (rewrite the
+ *     package directory) and read (`getModelFileText`); assert no
+ *     `ENOENT` is observed. On the pre-fix code, the read would fail
+ *     mid-rewrite. With the per-package mutex now covering both paths,
+ *     all reads succeed.
+ *
+ *  2. **Mutex coverage** — manually hold `withPackageLock` and assert
+ *     that a concurrent reader is pending until released. Pins the
+ *     invariant that readers actually take the lock.
+ *
+ *  3. **Download does not block compile** — start an `installPackage`
+ *     whose downloader never resolves on its own, then assert that
+ *     `getModelFileText` resolves promptly. This pins the Phase 1 /
+ *     Phase 2 split — if a future regression accidentally moves the
+ *     download inside the lock, this test fails.
+ */
+describe("package directory race", () => {
+   let rootDir: string;
+   let envPath: string;
+   let fixtureDir: string;
+
+   const PUBLISHER_JSON = JSON.stringify({
+      name: "pkg",
+      description: "race-test fixture",
+   });
+   const MODEL_MALLOY = `source: ones is duckdb.sql("SELECT 1 as x")\n`;
+
+   async function writeFixture(targetDir: string): Promise<void> {
+      await fs.mkdir(targetDir, { recursive: true });
+      await fs.writeFile(
+         path.join(targetDir, "publisher.json"),
+         PUBLISHER_JSON,
+      );
+      await fs.writeFile(path.join(targetDir, "model.malloy"), MODEL_MALLOY);
+   }
+
+   async function copyDir(src: string, dst: string): Promise<void> {
+      await fs.mkdir(dst, { recursive: true });
+      await fs.cp(src, dst, { recursive: true });
+   }
+
+   beforeEach(async () => {
+      rootDir = await fs.mkdtemp(path.join(os.tmpdir(), "publisher-race-"));
+      envPath = path.join(rootDir, "env");
+      fixtureDir = path.join(rootDir, "fixture");
+      await fs.mkdir(envPath, { recursive: true });
+      await writeFixture(fixtureDir);
+   });
+
+   afterEach(async () => {
+      await fs.rm(rootDir, { recursive: true, force: true }).catch(() => {});
+   });
+
+   it("(A) concurrent installs and reads never observe a half-rewritten tree", async () => {
+      const env = await Environment.create("testEnv", envPath, []);
+
+      // Initial install to populate the canonical path.
+      await env.installPackage("pkg", (stagingPath) =>
+         copyDir(fixtureDir, stagingPath),
+      );
+
+      const ITERATIONS = 30;
+      const errors: unknown[] = [];
+      let mutatorDone = false;
+
+      // Mutator loop: re-install the package over and over. Each iteration
+      // exercises the full Phase 1 (no-lock) + Phase 2 (locked) swap.
+      const mutator = (async () => {
+         try {
+            for (let i = 0; i < ITERATIONS; i++) {
+               try {
+                  await env.installPackage("pkg", (stagingPath) =>
+                     copyDir(fixtureDir, stagingPath),
+                  );
+               } catch (err) {
+                  errors.push({ kind: "install", err });
+               }
+            }
+         } finally {
+            mutatorDone = true;
+         }
+      })();
+
+      // Reader loop: hammer `getModelFileText` while installs run. On the
+      // pre-fix code (no lock on reads), the read would sometimes hit ENOENT
+      // because the canonical dir was momentarily missing during the rename
+      // window. With the per-package mutex covering reads as well, this
+      // window is never observable.
+      const reader = (async () => {
+         while (!mutatorDone) {
+            try {
+               const text = await env.getModelFileText("pkg", "model.malloy");
+               expect(text).toBe(MODEL_MALLOY);
+            } catch (err) {
+               errors.push({ kind: "read", err });
+            }
+         }
+      })();
+
+      await mutator;
+      await reader;
+
+      // Any error here means the lock wasn't actually covering one of the
+      // sides — that's the regression we're guarding against.
+      if (errors.length > 0) {
+         throw new Error(
+            `Observed ${errors.length} race-window error(s): ${JSON.stringify(
+               errors.slice(0, 3),
+               (_k, v) => (v instanceof Error ? `${v.name}: ${v.message}` : v),
+            )}`,
+         );
+      }
+   }, 60_000);
+
+   it("(B) compile-time disk reads queue behind withPackageLock", async () => {
+      const env = await Environment.create("testEnv", envPath, []);
+      await env.installPackage("pkg", (stagingPath) =>
+         copyDir(fixtureDir, stagingPath),
+      );
+
+      const lockEntered = defer<void>();
+      const releaseLock = defer<void>();
+
+      // Hold the per-package mutex from "outside" — simulates a mutator
+      // (install / delete / writePackageManifest) being in flight.
+      const lockHolder = env.withPackageLock("pkg", async () => {
+         lockEntered.resolve();
+         await releaseLock.promise;
+      });
+
+      await lockEntered.promise;
+
+      // While the lock is held, the reader must NOT make progress.
+      const readPromise = env.getModelFileText("pkg", "model.malloy");
+      const TIMEOUT_SENTINEL = Symbol("timeout");
+      const raced = await Promise.race([
+         readPromise,
+         new Promise<typeof TIMEOUT_SENTINEL>((resolve) =>
+            setTimeout(() => resolve(TIMEOUT_SENTINEL), 50),
+         ),
+      ]);
+      expect(raced).toBe(TIMEOUT_SENTINEL);
+
+      // Release the lock; the reader must now complete.
+      releaseLock.resolve();
+      await lockHolder;
+      const text = await readPromise;
+      expect(text).toBe(MODEL_MALLOY);
+   }, 15_000);
+
+   it("(C) a slow download does not block concurrent reads", async () => {
+      const env = await Environment.create("testEnv", envPath, []);
+      // Initial install to make the package present.
+      await env.installPackage("pkg", (stagingPath) =>
+         copyDir(fixtureDir, stagingPath),
+      );
+
+      const downloadGate = defer<void>();
+
+      // Kick off an install whose Phase 1 downloader stalls until we open
+      // the gate. Phase 2 (the brief locked swap) cannot run until then.
+      const slowInstall = env.installPackage("pkg", async (stagingPath) => {
+         await downloadGate.promise;
+         await copyDir(fixtureDir, stagingPath);
+      });
+
+      // The reader must resolve well before we open the gate, proving the
+      // per-package mutex is NOT held during Phase 1.
+      const readStart = Date.now();
+      const text = await env.getModelFileText("pkg", "model.malloy");
+      const readElapsedMs = Date.now() - readStart;
+
+      expect(text).toBe(MODEL_MALLOY);
+      // 1s is generous; in practice this resolves in single-digit ms.
+      expect(readElapsedMs).toBeLessThan(1_000);
+
+      // Now open the gate and let the install complete.
+      downloadGate.resolve();
+      await slowInstall;
+   }, 15_000);
+});
+
+interface Deferred<T> {
+   promise: Promise<T>;
+   resolve: (value: T) => void;
+   reject: (reason?: unknown) => void;
+}
+
+function defer<T>(): Deferred<T> {
+   let resolve!: (value: T) => void;
+   let reject!: (reason?: unknown) => void;
+   const promise = new Promise<T>((res, rej) => {
+      resolve = res;
+      reject = rej;
+   });
+   return { promise, resolve, reject };
+}