Pool host-approved restricted code with scoped and fair graph execution

restricted-execution/service.mjs

@@ -14,6 +14,46 @@ const absolute = (value) =>
   value.startsWith("/") &&
   !/[\s\u0000-\u001f]/.test(value) &&
   !value.split("/").includes("..");
+
+/** Host-wide engine reservations. Only trusted adapters hold these callbacks;
+ * workers cannot alter admission or evict another invocation's engine. */
+export class RestrictedEngineBudget {
+  #leases = new Set();
+  constructor(limit = 8, reserve = async () => ({})) {
+    if (!Number.isSafeInteger(limit) || limit < 1 || limit > 64) throw fixedError("INVALID_ENGINE_BUDGET");
+    this.limit = limit;
+    this.reserve = reserve;
+  }
+  async acquire(idle, evict) {
+    while (this.#leases.size >= this.limit) {
+      const victim = [...this.#leases].find((entry) => entry.idle());
+      if (!victim) throw fixedError("RESTRICTED_ENGINE_LIMIT");
+      await victim.evict();
+    }
+    const lease = { idle, evict };
+    this.#leases.add(lease);
+    try {
+      const resource = await this.reserve();
+      let releasing;
+      const release = () =>
+        (releasing ??= (async () => {
+          try {
+            await resource.release?.();
+          } finally {
+            this.#leases.delete(lease);
+          }
+        })());
+      release.attach = resource.attach;
+      return release;
+    } catch (error) {
+      this.#leases.delete(lease);
+      throw error;
+    }
+  }
+  stats() {
+    return { engines: this.#leases.size, limit: this.limit };
+  }
+}
 async function boundedSource(path, maximum, code) {
   const handle = await open(path, constants.O_RDONLY | constants.O_NOFOLLOW | constants.O_NONBLOCK);
   try {
@@ -101,12 +141,14 @@ export async function serveRestrictedRegistry(config, io = process) {
   const sdk = [];
   for (const m of config.sdk)
     sdk.push({ ...m, source: await boundedSource(m.path, 2 * 1024 * 1024, "RESTRICTED_SDK_LIMIT") });
+  const engineBudget = config.engineBudget ?? new RestrictedEngineBudget(limits.maxActive + limits.maxNested);
   const active = new Map(),
     engines = new Set(),
     idleEngines = [];
   const metrics = { starts: 0, reused: 0, completed: 0, peakRSSBytes: 0 };
   let buffered = Buffer.alloc(0),
-    stopped = false;
+    stopped = false,
+    failure;
   const queued = [];
   const send = (value) => {
     if (stopped) return;
@@ -191,15 +233,17 @@ export async function serveRestrictedRegistry(config, io = process) {
     })();
     return call.retiring;
   };
-  const fail = async (error) => {
+  const fail = (error) => {
-    if (stopped) return;
+    if (failure) return failure;
     stopped = true;
+    return (failure = (async () => {
       io.stdin.destroy();
       for (const call of [...queued]) await retire(call, error);
       await Promise.all([...active.values()].map((call) => retire(call, error)));
       await Promise.all([...engines].map(destroyEngine));
       io.stdout.end();
       io.exitCode = 1;
+    })());
   };
   const body = async (stream) => {
     const chunks = [];
@@ -228,8 +272,12 @@ export async function serveRestrictedRegistry(config, io = process) {
       }
       engine.server?.closeAllConnections();
       if (engine.server) await new Promise((resolve) => engine.server.close(resolve));
+      try {
         if (engine.directory) await rm(engine.directory, { recursive: true, force: true });
+      } finally {
         engines.delete(engine);
+        await engine.release?.();
+      }
     })();
     return engine.destroying;
   };
@@ -248,6 +296,11 @@ export async function serveRestrictedRegistry(config, io = process) {
     const engine = { current: call, cpuTotal: 0 };
     engines.add(engine);
     call.engine = engine;
+    engine.release = await engineBudget.acquire(
+      () => !engine.current && !engine.destroying,
+      () => destroyEngine(engine),
+    );
+    if (!call.open) return;
     engine.directory = await mkdtemp(join(tmpdir(), "qx-engine-"));
     if (!call.open) return;
     const broker = join(engine.directory, "broker.sock");
@@ -339,6 +392,7 @@ export async function serveRestrictedRegistry(config, io = process) {
       engine.child.once("error", resolve);
     });
     call.child = engine.child;
+    await engine.release.attach?.(engine.child.pid);
     watchCPU(call);
     let diagnostics = 0;
     engine.child.stderr.on("data", (chunk) => {

restricted-execution/service.test.mjs

@@ -5,7 +5,7 @@ import { mkdtemp, writeFile, rm } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { fileURLToPath } from "node:url";
-import { serveRestrictedRegistry } from "./service.mjs";
+import { serveRestrictedRegistry, RestrictedEngineBudget } from "./service.mjs";
 
 const workerd = process.env.QUIXOS_TEST_WORKERD;
 if (!workerd) throw Error("Set QUIXOS_TEST_WORKERD to the pinned engine");
@@ -571,3 +571,28 @@ test("detached digest callbacks cannot use a completed frame or contaminate reus
     await f.close();
   }
 });
+
+test("the shared engine budget evicts idle code families and rejects busy oversubscription", async () => {
+  const engineBudget = new RestrictedEngineBudget(1);
+  const a = await fixture(registry("return request.input;"), { engineBudget });
+  const b = await fixture(
+    registry("return await channel({portId:'read',action:{kind:'state',primitive:'read'},input:null});"),
+    { engineBudget },
+  );
+  try {
+    a.start("first family");
+    assert.equal((await a.receive()).output, "first family");
+    const busy = b.start();
+    const port = await b.receive();
+    assert.equal(a.stats().engines, 0, "Idle foreign code is evicted before allocating another engine");
+    assert.equal(engineBudget.stats().engines, 1);
+    a.start("overload");
+    assert.equal((await a.receive()).error.code, "RESTRICTED_ENGINE_LIMIT");
+    assert.equal(engineBudget.stats().engines, 1);
+    b.send({ kind: "port-result", ...busy, sequence: port.sequence, status: "returned", output: "survived" });
+    assert.equal((await b.receive()).output, "survived");
+  } finally {
+    await Promise.all([a.close(), b.close()]);
+  }
+  assert.equal(engineBudget.stats().engines, 0);
+});