@malloy-publisher/server 0.0.197 → 0.0.198

package/src/service/environment_admission.spec.ts

@@ -0,0 +1,180 @@
+import { afterEach, beforeEach, describe, expect, it } from "bun:test";
+import * as fs from "fs";
+import * as os from "os";
+import * as path from "path";
+
+import { ServiceUnavailableError } from "../errors";
+import { buildEnvironmentMalloyConfig } from "./connection";
+import { Environment } from "./environment";
+import type { PackageMemoryGovernor } from "./package_memory_governor";
+
+/**
+ * Minimal subset of {@link PackageMemoryGovernor} that
+ * `Environment.assertCanAdmitNewPackage` actually consults. Allows us
+ * to drive the gate from the tests without spinning the real OTel
+ * instrumentation pipeline.
+ */
+class StubGovernor {
+   public backpressured = false;
+   isBackpressured(): boolean {
+      return this.backpressured;
+   }
+}
+
+function makeEnvironment(envPath: string): Environment {
+   const malloyConfig = buildEnvironmentMalloyConfig([], envPath);
+   return new Environment("test-env", envPath, malloyConfig, []);
+}
+
+describe("Environment admission gate (memory governor choke point)", () => {
+   let envDir: string;
+
+   beforeEach(() => {
+      envDir = fs.mkdtempSync(
+         path.join(os.tmpdir(), "publisher-env-admission-"),
+      );
+   });
+
+   afterEach(() => {
+      fs.rmSync(envDir, { recursive: true, force: true });
+   });
+
+   it("admits new packages when no governor is attached (legacy behaviour)", async () => {
+      const env = makeEnvironment(envDir);
+      // No governor set; the gate must be a pure no-op. The package
+      // doesn't exist on disk, so addPackage rejects with
+      // PackageNotFoundError — that we get any error other than 503
+      // is the assertion.
+      let caught: unknown;
+      try {
+         await env.addPackage("does-not-exist");
+      } catch (err) {
+         caught = err;
+      }
+      expect(caught).toBeDefined();
+      expect(caught).not.toBeInstanceOf(ServiceUnavailableError);
+   });
+
+   it("rejects getPackage cache-miss with 503 when back-pressured", async () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      // No package by this name has ever been loaded, so this is the
+      // exact "lazy-load on cache miss" path Monty flagged. The gate
+      // must throw before Package.create touches the disk.
+      await expect(env.getPackage("ghost", false)).rejects.toBeInstanceOf(
+         ServiceUnavailableError,
+      );
+   });
+
+   it("rejects getPackage reload=true with 503 when back-pressured", async () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      await expect(env.getPackage("ghost", true)).rejects.toBeInstanceOf(
+         ServiceUnavailableError,
+      );
+   });
+
+   it("rejects addPackage with 503 when back-pressured (after the 404 check passes)", async () => {
+      // Create a real (empty) package directory so the existence
+      // check passes and the gate gets to run. Without this, the
+      // PackageNotFoundError would mask the 503 we want to assert.
+      const pkgName = "real-pkg";
+      fs.mkdirSync(path.join(envDir, pkgName));
+
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      await expect(env.addPackage(pkgName)).rejects.toBeInstanceOf(
+         ServiceUnavailableError,
+      );
+   });
+
+   it("returns 404 (not 503) when the package directory does not exist, even under pressure", async () => {
+      // 404 must take precedence over 503: a permanent "you forgot to
+      // upload the package" error should not be masked as a transient
+      // "retry later" — otherwise operators chase phantom memory
+      // problems while the real fix is a missing artifact.
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      let caught: unknown;
+      try {
+         await env.addPackage("never-existed");
+      } catch (err) {
+         caught = err;
+      }
+      expect(caught).toBeDefined();
+      expect(caught).not.toBeInstanceOf(ServiceUnavailableError);
+   });
+
+   it("allowAdmission=true bypasses the gate (for future warmup/probe callers)", async () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      // With the bypass, the gate must not fire — the call should
+      // proceed to the real loader and fail there with some other
+      // error (PackageNotFoundError-equivalent from Package.create on
+      // a non-existent directory). The assertion is "not 503".
+      let caught: unknown;
+      try {
+         await env.getPackage("ghost", false, { allowAdmission: true });
+      } catch (err) {
+         caught = err;
+      }
+      expect(caught).toBeDefined();
+      expect(caught).not.toBeInstanceOf(ServiceUnavailableError);
+   });
+
+   it("clearing back-pressure on the governor immediately re-admits new loads", async () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+
+      governor.backpressured = true;
+      await expect(env.getPackage("ghost", false)).rejects.toBeInstanceOf(
+         ServiceUnavailableError,
+      );
+
+      // Flip the flag (simulating the periodic poller crossing the
+      // low-water mark) and verify the next call no longer 503s.
+      governor.backpressured = false;
+      let caught: unknown;
+      try {
+         await env.getPackage("ghost", false);
+      } catch (err) {
+         caught = err;
+      }
+      expect(caught).toBeDefined();
+      expect(caught).not.toBeInstanceOf(ServiceUnavailableError);
+   });
+
+   it("detaching the governor (set null) reverts to legacy admit-everything", async () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      env.setMemoryGovernor(null);
+
+      let caught: unknown;
+      try {
+         await env.getPackage("ghost", false);
+      } catch (err) {
+         caught = err;
+      }
+      expect(caught).toBeDefined();
+      expect(caught).not.toBeInstanceOf(ServiceUnavailableError);
+   });
+});

package/src/service/environment_store.spec.ts

@@ -355,15 +355,6 @@ describe("EnvironmentStore Service", () => {
       expect(projects.length).toBe(2);
       expect(projects.map((p) => p.name)).toContain(projectName1);
       expect(projects.map((p) => p.name)).toContain(projectName2);
-
-      // All envs initialized cleanly → status is "serving" (not
-      // "degraded") and there's no failedEnvironments key on the
-      // response. This is the happy-path companion to the
-      // "should skip a project with invalid startup connection config"
-      // test which exercises the degraded path.
-      const status = await newEnvironmentStore.getStatus();
-      expect(status.operationalState).toBe("serving");
-      expect(status.failedEnvironments).toBeUndefined();
    });
 
    it("should skip a project with invalid startup connection config", async () => {
@@ -436,16 +427,6 @@ describe("EnvironmentStore Service", () => {
       await expect(
          newEnvironmentStore.getEnvironment(invalidProjectName),
       ).rejects.toThrow();
-
-      // The skipped environment should surface in the status response
-      // so external callers (CI smoke tests, dashboards) can tell the
-      // server is only partially serving.
-      const status = await newEnvironmentStore.getStatus();
-      expect(status.operationalState).toBe("degraded");
-      expect(status.failedEnvironments).toBeDefined();
-      expect(status.failedEnvironments?.map((f) => f.name)).toContain(
-         invalidProjectName,
-      );
    });
 
    it("should handle project updates", async () => {

package/src/service/environment_store.ts

@@ -25,16 +25,12 @@ import {
    FrozenConfigError,
    PackageNotFoundError,
 } from "../errors";
-import {
-   getOperationalState,
-   markDegraded,
-   markNotReady,
-   markReady,
-} from "../health";
+import { getOperationalState, markNotReady, markReady } from "../health";
 import { formatDuration, logger } from "../logger";
 import { Connection } from "../storage/DatabaseInterface";
 import { StorageConfig, StorageManager } from "../storage/StorageManager";
 import { Environment, PackageStatus } from "./environment";
+import type { PackageMemoryGovernor } from "./package_memory_governor";
 type ApiEnvironment = components["schemas"]["Environment"];
 
 const AZURE_SUPPORTED_SCHEMES = ["https://", "http://", "abfss://", "az://"];
@@ -101,12 +97,15 @@ export class EnvironmentStore {
    public publisherConfigIsFrozen: boolean;
    public finishedInitialization: Promise<void>;
    private isInitialized: boolean = false;
-   private failedEnvironments: Array<{ name: string; error: string }> = [];
    public storageManager: StorageManager;
    private s3Client = new S3({
       followRegionRedirects: true,
    });
    private gcsClient: Storage;
+   // Shared by every Environment so the back-pressure decision is
+   // process-wide. Set once at server start via setMemoryGovernor;
+   // new Environments pick it up at construction.
+   private memoryGovernor: PackageMemoryGovernor | null = null;
 
    constructor(serverRootPath: string) {
       this.serverRootPath = serverRootPath;
@@ -123,6 +122,19 @@ export class EnvironmentStore {
       this.finishedInitialization = this.initialize();
    }
 
+   /**
+    * Attach (or detach with `null`) the shared {@link PackageMemoryGovernor}.
+    * Propagated to every Environment so the back-pressure decision is
+    * process-wide, and remembered so any Environment created *after*
+    * this call also picks it up at construction.
+    */
+   public setMemoryGovernor(governor: PackageMemoryGovernor | null): void {
+      this.memoryGovernor = governor;
+      for (const env of this.environments.values()) {
+         env.setMemoryGovernor(governor);
+      }
+   }
+
    private async addConfiguredEnvironment(environment: ProcessedEnvironment) {
       try {
          await this.addEnvironment(
@@ -148,10 +160,6 @@ export class EnvironmentStore {
          `Error initializing environment${label}; skipping environment`,
          this.extractErrorDataFromError(error),
       );
-      this.failedEnvironments.push({
-         name: environmentName ?? "<unknown>",
-         error: error instanceof Error ? error.message : String(error),
-      });
    }
 
    private async initialize() {
@@ -248,6 +256,9 @@ export class EnvironmentStore {
                               ...conn.config,
                            })),
                         );
+                        environmentInstance.setMemoryGovernor(
+                           this.memoryGovernor,
+                        );
 
                         // Get packages from database
                         const packages = await repository.listPackages(
@@ -285,11 +296,7 @@ export class EnvironmentStore {
          }
 
          this.isInitialized = true;
-         if (this.failedEnvironments.length > 0) {
-            markDegraded();
-         } else {
-            markReady();
-         }
+         markReady();
          const initializationDuration = performance.now() - initialTime;
          logger.info(
             `Environment store successfully initialized in ${formatDuration(initializationDuration)}`,
@@ -703,11 +710,6 @@ export class EnvironmentStore {
          frozenConfig: isPublisherConfigFrozen(this.serverRootPath),
          operationalState:
             getOperationalState() as components["schemas"]["ServerStatus"]["operationalState"],
-         ...(this.failedEnvironments.length > 0 && {
-            failedEnvironments: [
-               ...this.failedEnvironments,
-            ] as components["schemas"]["ServerStatus"]["failedEnvironments"],
-         }),
       };
 
       const environments = await this.listEnvironments(true);
@@ -856,6 +858,7 @@ export class EnvironmentStore {
          absoluteEnvironmentPath,
          environment.connections || [],
       );
+      newEnvironment.setMemoryGovernor(this.memoryGovernor);
 
       if (!newEnvironment.metadata) newEnvironment.metadata = {};
       newEnvironment.metadata.location = absoluteEnvironmentPath;

package/src/service/manifest_service.spec.ts

@@ -39,6 +39,7 @@ function createMocks() {
    } as unknown as sinon.SinonStubbedInstance<ResourceRepository>;
 
    const reloadAllModels = sandbox.stub().resolves();
+   const reloadAllModelsForPackage = sandbox.stub().resolves();
 
    const pkg = {
       reloadAllModels,
@@ -46,6 +47,7 @@ function createMocks() {
 
    const environment = {
       getPackage: sandbox.stub().resolves(pkg),
+      reloadAllModelsForPackage,
    };
 
    const environmentStore = {
@@ -66,6 +68,7 @@ function createMocks() {
       environment,
       pkg,
       reloadAllModels,
+      reloadAllModelsForPackage,
       service,
    };
 }
@@ -153,7 +156,9 @@ describe("ManifestService", () => {
             ),
          ).toBe(true);
          expect(ctx.environment.getPackage.calledWith("pkg", false)).toBe(true);
-         expect(ctx.reloadAllModels.calledWith(manifest.entries)).toBe(true);
+         expect(
+            ctx.reloadAllModelsForPackage.calledWith("pkg", manifest.entries),
+         ).toBe(true);
       });
 
       it("should return an empty manifest when no entries exist", async () => {
@@ -170,7 +175,7 @@ describe("ManifestService", () => {
          );
 
          expect(result.entries).toEqual({});
-         expect(ctx.reloadAllModels.calledWith({})).toBe(true);
+         expect(ctx.reloadAllModelsForPackage.calledWith("pkg", {})).toBe(true);
       });
    });
 

package/src/service/manifest_service.ts

@@ -82,8 +82,14 @@ export class ManifestService {
          environmentName,
          false,
       );
-      const pkg = await environment.getPackage(packageName, false);
-      await pkg.reloadAllModels(manifest.entries);
+      // Ensure the package is loaded, then reload its models under the
+      // per-package mutex so the disk reads are serialized against
+      // installPackage / deletePackage.
+      await environment.getPackage(packageName, false);
+      await environment.reloadAllModelsForPackage(
+         packageName,
+         manifest.entries,
+      );
 
       logger.info("Reloaded manifest and recompiled models", {
          environmentId,

package/src/service/materialization_service.ts

@@ -376,8 +376,14 @@ export class MaterializationService {
                environmentName,
                false,
             );
-            const pkg = await environment.getPackage(packageName, false);
-            await pkg.reloadAllModels(updatedManifest.entries);
+            // Ensure the package is loaded, then reload models under the
+            // per-package mutex so the disk reads are serialized against
+            // installPackage / deletePackage.
+            await environment.getPackage(packageName, false);
+            await environment.reloadAllModelsForPackage(
+               packageName,
+               updatedManifest.entries,
+            );
          }
 
          await this.transitionExecution(executionId, "SUCCESS", {
@@ -604,8 +610,13 @@ export class MaterializationService {
       // ── STEP 2: COMPILE & PLAN ─────────────────────────────────────
       // `connections` is built lazily from the connection names the plan
       // actually targets — no upfront ATTACH on every environment connection.
+      // Hold the per-package mutex for the duration of the compile so the
+      // `fs.stat` + `runtime.loadModel` calls inside `compilePackageBuildPlan`
+      // are serialized against `installPackage` / `deletePackage`.
       const { graphs, sources, connectionDigests, connections } =
-         await this.compilePackageBuildPlan(pkg, signal);
+         await environment.withPackageLock(packageName, () =>
+            this.compilePackageBuildPlan(pkg, signal),
+         );
 
       if (graphs.length === 0) {
          logger.info("No persist sources to build");

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");
+   });
+});