quixos-toolchain/quixos-nix-helpers
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Reuse restricted engines with fresh invocation state and revocable ports

Browse Source
This commit is contained in:
Timothy J. Aveni
2026-10-07 16:14:15 -07:00
parent f934347041
commit 929fdb4a28
3 changed files with 377 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
restricted-execution/service.mjs
+216 -110
View File
@@ -1,4 +1,5 @@
import { spawn } from "node:child_process";
import { randomUUID } from "node:crypto";
import { createServer, request as httpRequest } from "node:http";
import { mkdtemp, readFile, writeFile, rm, open } from "node:fs/promises";
import { constants } from "node:fs";
@@ -34,7 +35,8 @@ async function boundedSource(path, maximum, code) {
/** Trusted stdio adapter inside one workspace VM. Engine processes never own
* coordinator credentials. The supervising service must impose VM/OS limits.
* Per-frame process replacement prevents retained globals/handles crossing calls. */
* Each engine runs one fresh Worker per invocation. Host tokens expire before
* reuse; cancellation and watchdogs destroy the affected engine, not neighbors. */
export async function serveRestrictedRegistry(config, io = process) {
if (
!object(config) ||
@@ -68,6 +70,7 @@ export async function serveRestrictedRegistry(config, io = process) {
maxQueued: config.maxQueued ?? 8,
maxActive: config.maxActive ?? 2,
maxNested: config.maxNested ?? 4,
engineIdleMs: config.engineIdleMs ?? 30000,
};
if (
!Number.isSafeInteger(limits.maxBytes) ||
@@ -87,7 +90,10 @@ export async function serveRestrictedRegistry(config, io = process) {
limits.maxActive > 16 ||
!Number.isSafeInteger(limits.maxNested) ||
limits.maxNested < 1 ||
limits.maxNested > 16
limits.maxNested > 16 ||
!Number.isSafeInteger(limits.engineIdleMs) ||
limits.engineIdleMs < 1 ||
limits.engineIdleMs > 60000
)
throw fixedError("INVALID_RESTRICTED_LIMITS");
const packageSource = await boundedSource(config.module, 1024 * 1024, "RESTRICTED_SOURCE_LIMIT");
@@ -95,7 +101,10 @@ 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 active = new Map();
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;
const queued = [];
@@ -121,7 +130,9 @@ export async function serveRestrictedRegistry(config, io = process) {
throw error;
}
const fields = text.slice(text.lastIndexOf(")") + 2).split(" ");
const cpuMicroseconds = (Number(fields[11]) + Number(fields[12])) * 10000;
const total = (Number(fields[11]) + Number(fields[12])) * 10000;
const cpuMicroseconds = total - (call.cpuBaseline ?? 0);
if (call.engine) call.engine.cpuTotal = total;
if (!Number.isSafeInteger(cpuMicroseconds) || cpuMicroseconds < (call.cpuMicroseconds ?? 0))
throw fixedError("RESTRICTED_CPU_ACCOUNTING");
if (cpuMicroseconds !== (call.cpuMicroseconds ?? 0)) {
@@ -158,17 +169,23 @@ export async function serveRestrictedRegistry(config, io = process) {
} catch (error) {
usageFailure = error;
}
if (call.child?.pid) {
try {
process.kill(-call.child.pid, "SIGKILL");
} catch (error) {
if (error.code !== "ESRCH") throw error;
}
await call.exited;
const engine = call.engine;
if (engine) {
// Revoke the frame token before returning an engine to the idle pool.
// A delayed child request can never acquire the next call's authority.
if (engine.current === call) engine.current = null;
if (!stopped && reason.code === "INVOCATION_ENDED" && !usageFailure && !engine.destroying) {
const status = await readFile(`/proc/${engine.child.pid}/status`, "utf8").catch(() => "");
const rss = Number(status.match(/^VmRSS:\s+(\d+)/m)?.[1] ?? Infinity) * 1024;
metrics.peakRSSBytes = Math.max(metrics.peakRSSBytes, rss);
// Bound retained engine heaps as well as lifetime. Hard aggregate
// physical memory remains the supervising process/cgroup's contract.
if (rss <= 128 * 1024 * 1024) {
idleEngines.push(engine);
engine.idleTimer = setTimeout(() => void destroyEngine(engine).catch(fail), limits.engineIdleMs);
} else await destroyEngine(engine);
} else await destroyEngine(engine);
}
call.server?.closeAllConnections();
if (call.server) await new Promise((resolve) => call.server.close(resolve));
if (call.directory) await rm(call.directory, { recursive: true, force: true });
active.delete(call.message.invocation);
if (usageFailure) throw usageFailure;
})();
@@ -180,6 +197,7 @@ export async function serveRestrictedRegistry(config, io = process) {
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;
};
@@ -193,6 +211,148 @@ export async function serveRestrictedRegistry(config, io = process) {
}
return JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(Buffer.concat(chunks)));
};
const destroyEngine = async (engine) => {
if (engine.destroying) return engine.destroying;
engine.destroying = (async () => {
clearTimeout(engine.idleTimer);
const index = idleEngines.indexOf(engine);
if (index >= 0) idleEngines.splice(index, 1);
engine.current = null;
if (engine.child?.pid) {
try {
process.kill(-engine.child.pid, "SIGKILL");
} catch (error) {
if (error.code !== "ESRCH") throw error;
}
await engine.exited;
}
engine.server?.closeAllConnections();
if (engine.server) await new Promise((resolve) => engine.server.close(resolve));
if (engine.directory) await rm(engine.directory, { recursive: true, force: true });
engines.delete(engine);
})();
return engine.destroying;
};
const prepareEngine = async (call) => {
const cached = idleEngines.pop();
if (cached) {
clearTimeout(cached.idleTimer);
cached.current = call;
call.engine = cached;
call.child = cached.child;
call.cpuBaseline = cached.cpuTotal ?? 0;
metrics.reused++;
watchCPU(call);
return;
}
const engine = { current: call, cpuTotal: 0 };
engines.add(engine);
call.engine = engine;
engine.directory = await mkdtemp(join(tmpdir(), "qx-engine-"));
if (!call.open) return;
const broker = join(engine.directory, "broker.sock");
engine.socket = join(engine.directory, "engine.sock");
engine.server = createServer(async (request, response) => {
const call = engine.current;
const message = call?.message;
try {
if (!call?.open || request.method !== "POST" || request.url !== "/port/" + call.token)
throw fixedError("EXPIRED_CONTEXT");
const port = await body(request);
if (
!call.open ||
!object(port) ||
Object.keys(port).some((key) => !["portId", "action", "input", "key"].includes(key)) ||
typeof port.portId !== "string" ||
!object(port.action) ||
!Object.hasOwn(port, "input")
)
throw fixedError("INVALID_RESTRICTED_PORT");
if (
++call.sequence > Math.min(limits.maxPortCalls, message.budget?.calls ?? limits.maxPortCalls) ||
call.pending.size >= limits.maxPortCalls
)
throw fixedError("RESTRICTED_PORT_LIMIT");
const sequence = call.sequence;
const result = new Promise((resolve, reject) => call.pending.set(sequence, { resolve, reject }));
void result.catch(() => {});
send({
kind: "port",
invocation: message.invocation,
frame: message.frame,
sequence,
portId: port.portId,
action: { ...port.action, ...(port.key === undefined ? {} : { key: port.key }) },
input: port.input,
});
const value = await result;
if (!call.open) throw fixedError("EXPIRED_CONTEXT");
const encoded = JSON.stringify(value);
if (Buffer.byteLength(encoded) > limits.maxBytes) throw fixedError("RESTRICTED_MESSAGE_LIMIT");
response.writeHead(200, { "content-type": "application/json" });
response.end(encoded);
} catch {
response.writeHead(409);
response.end('{"status":"failed","error":{"code":"RESTRICTED_PORT_FAILED"}}');
}
});
await new Promise((resolve, reject) => {
engine.server.once("error", reject);
engine.server.listen(broker, resolve);
});
const modules = { "entry.mjs": { js: workerSource }, "package.mjs": { js: packageSource } };
for (const module of sdk) modules[module.name] = { js: module.source };
const setting = { registryExport: config.registryExport, exports: config.exports, maxBytes: limits.maxBytes };
// The parent contains immutable checked source only. It never imports the
// authored module or keeps a child stub, request, result or frame handle.
const parent = `const modules=${JSON.stringify(modules)}, setting=${JSON.stringify(setting)};
export default {async fetch(request,env) {
const input=await request.json();
const child=env.LOADER.load({compatibilityDate:"2026-10-05",mainModule:"entry.mjs",modules,
globalOutbound:null,env:{BROKER:env.BROKER,CONFIG:{...setting,frameToken:input.frameToken}}});
const response=await child.getEntrypoint().fetch(new Request("http://invocation/",{
method:"POST",body:JSON.stringify(input.message)}));
const body=await response.arrayBuffer();
return new Response(body,{status:response.status,headers:response.headers});
}};`;
await writeFile(join(engine.directory, "parent.mjs"), parent);
const configuration = `using Workerd = import "/workerd/workerd.capnp"; const config :Workerd.Config = (
services=[(name="main",worker=(modules=[(name="parent.mjs",esModule=embed "parent.mjs")],
compatibilityDate="2026-10-05",compatibilityFlags=["experimental"],globalOutbound="deny",
bindings=[(name="LOADER",workerLoader=()),(name="BROKER",service="broker")])),
(name="broker",external=(address=${JSON.stringify("unix:" + broker)},http=())),
(name="deny",network=(allow=[]))],
sockets=[(name="frame",address=${JSON.stringify("unix:" + engine.socket)},http=(),service="main")]);`;
const file = join(engine.directory, "config.capnp");
await writeFile(file, configuration);
if (!call.open) return;
engine.child = spawn(config.workerd, ["serve", "--experimental", file], {
detached: true,
env: { TZ: "UTC", LANG: "C.UTF-8" },
cwd: engine.directory,
stdio: ["ignore", "ignore", "pipe"],
});
metrics.starts++;
engine.exited = new Promise((resolve) => {
engine.child.once("exit", resolve);
engine.child.once("error", resolve);
});
call.child = engine.child;
watchCPU(call);
let diagnostics = 0;
engine.child.stderr.on("data", (chunk) => {
if ((diagnostics += chunk.length) > 65536) {
if (engine.current?.open) engine.current.terminate("RESTRICTED_DIAGNOSTIC_LIMIT");
else void destroyEngine(engine).catch(fail);
}
});
engine.child.once("error", () => engine.current?.terminate("RESTRICTED_START_FAILED"));
void engine.exited.then(() => {
if (engine.current?.open) engine.current.terminate("RESTRICTED_ENGINE_EXITED");
else void destroyEngine(engine).catch(fail);
});
};
const available = (call) =>
[...active.values()].filter((entry) => entry.nested === call.nested).length <
(call.nested ? limits.maxNested : limits.maxActive);
@@ -207,12 +367,21 @@ export async function serveRestrictedRegistry(config, io = process) {
if (reply.kind === "cancelled") call.cancellationReply = reply;
return (call.finished ??= (async () => {
await retire(call, reason);
metrics.lastInvocation = {
queueMs: (call.startedAt ?? performance.now()) - call.receivedAt,
preparationMs: call.preparationMs ?? 0,
roundtripMs: call.roundtripMs ?? 0,
initializationMs: call.initializationMs ?? null,
executionMs: call.executionMs ?? null,
totalMs: performance.now() - call.receivedAt,
};
send(call.cancellationReply ?? reply);
pump();
})());
};
const invoke = async (call) => {
const { message } = call;
call.startedAt = performance.now();
active.set(message.invocation, call);
const terminate = (code) => {
if (!call.open) return;
@@ -224,100 +393,12 @@ export async function serveRestrictedRegistry(config, io = process) {
};
call.terminate = terminate;
try {
let engine;
call.preparing = (async () => {
call.directory = await mkdtemp(join(tmpdir(), "qx-frame-"));
if (!call.open) return;
const broker = join(call.directory, "broker.sock");
engine = join(call.directory, "engine.sock");
call.server = createServer(async (request, response) => {
try {
if (!call.open || request.method !== "POST" || request.url !== "/port") throw fixedError("EXPIRED_CONTEXT");
const port = await body(request);
if (
!call.open ||
!object(port) ||
Object.keys(port).some((key) => !["portId", "action", "input", "key"].includes(key)) ||
typeof port.portId !== "string" ||
!object(port.action) ||
!Object.hasOwn(port, "input")
)
throw fixedError("INVALID_RESTRICTED_PORT");
if (
++call.sequence > Math.min(limits.maxPortCalls, message.budget?.calls ?? limits.maxPortCalls) ||
call.pending.size >= limits.maxPortCalls
)
throw fixedError("RESTRICTED_PORT_LIMIT");
const sequence = call.sequence;
const result = new Promise((resolve, reject) => call.pending.set(sequence, { resolve, reject }));
void result.catch(() => {});
send({
kind: "port",
invocation: message.invocation,
frame: message.frame,
sequence,
portId: port.portId,
action: { ...port.action, ...(port.key === undefined ? {} : { key: port.key }) },
input: port.input,
});
const value = await result;
if (!call.open) throw fixedError("EXPIRED_CONTEXT");
const encoded = JSON.stringify(value);
if (Buffer.byteLength(encoded) > limits.maxBytes) throw fixedError("RESTRICTED_MESSAGE_LIMIT");
response.writeHead(200, { "content-type": "application/json" });
response.end(encoded);
} catch {
response.writeHead(409);
response.end('{"status":"failed","error":{"code":"RESTRICTED_PORT_FAILED"}}');
}
});
await new Promise((resolve, reject) => {
call.server.once("error", reject);
call.server.listen(broker, resolve);
});
await writeFile(join(call.directory, "package.mjs"), packageSource);
await writeFile(join(call.directory, "entry.mjs"), workerSource);
const modules = [
{ name: "entry.mjs", esModule: 'embed "entry.mjs"' },
{ name: "package.mjs", esModule: 'embed "package.mjs"' },
];
for (const [index, module] of sdk.entries()) {
const file = `sdk-${index}.mjs`;
await writeFile(join(call.directory, file), module.source);
modules.push({ name: module.name, esModule: "embed " + JSON.stringify(file) });
}
const setting = { registryExport: config.registryExport, exports: config.exports, maxBytes: limits.maxBytes };
const configuration = `using Workerd = import "/workerd/workerd.capnp"; const config :Workerd.Config = (
services = [(name="main", worker=(modules=[${modules.map((m) => `(name=${JSON.stringify(m.name)}, esModule=${m.esModule})`).join(",")}],
compatibilityDate="2026-10-05", globalOutbound="deny",
bindings=[(name="BROKER", service="broker"),(name="CONFIG",json=${JSON.stringify(JSON.stringify(setting))})])),
(name="broker",external=(address=${JSON.stringify("unix:" + broker)},http=())), (name="deny",network=(allow=[]))],
sockets=[(name="frame",address=${JSON.stringify("unix:" + engine)},http=(),service="main")]);`;
const file = join(call.directory, "config.capnp");
await writeFile(file, configuration);
if (!call.open) return;
call.child = spawn(config.workerd, ["serve", file], {
detached: true,
env: { TZ: "UTC", LANG: "C.UTF-8" },
cwd: call.directory,
stdio: ["ignore", "ignore", "pipe"],
});
call.exited = new Promise((resolve) => {
call.child.once("exit", resolve);
call.child.once("error", resolve);
});
watchCPU(call);
let diagnostics = 0;
call.child.stderr.on("data", (chunk) => {
if ((diagnostics += chunk.length) > 65536) terminate("RESTRICTED_DIAGNOSTIC_LIMIT");
});
call.child.once("error", () => terminate("RESTRICTED_START_FAILED"));
void call.exited.then(() => {
if (call.open) terminate("RESTRICTED_ENGINE_EXITED");
});
})();
call.token = randomUUID();
call.preparing = prepareEngine(call);
await call.preparing;
call.preparationMs = performance.now() - call.startedAt;
if (!call.open) return;
const roundtripAt = performance.now();
const result = await new Promise((resolve, reject) => {
const attempt = () => {
if (!call.open) {
@@ -325,10 +406,17 @@ export async function serveRestrictedRegistry(config, io = process) {
return;
}
const request = httpRequest(
{ socketPath: engine, path: "/", method: "POST", agent: false },
{ socketPath: call.engine.socket, path: "/", method: "POST", agent: false },
async (response) => {
try {
const result = await body(response);
for (const [header, key] of [
["x-quixos-initialization-ms", "initializationMs"],
["x-quixos-execution-ms", "executionMs"],
]) {
const value = Number(response.headers[header]);
if (Number.isFinite(value) && value >= 0 && value <= limits.deadlineMs) call[key] = value;
}
if (response.statusCode !== 200) throw fixedError("RESTRICTED_PACKAGE_FAILED");
resolve(result);
} catch (error) {
@@ -340,10 +428,11 @@ export async function serveRestrictedRegistry(config, io = process) {
if (call.open && ["ENOENT", "ECONNREFUSED"].includes(error.code)) setTimeout(attempt, 5);
else reject(error);
});
request.end(JSON.stringify(message));
request.end(JSON.stringify({ message, frameToken: call.token }));
};
attempt();
});
call.roundtripMs = performance.now() - roundtripAt;
if (!call.open) return;
if (
!object(result) ||
@@ -367,6 +456,7 @@ export async function serveRestrictedRegistry(config, io = process) {
error: { code: result.error.code, message: result.error.code },
};
await finish(call, fixedError("INVOCATION_ENDED"), reply);
metrics.completed++;
} catch (error) {
if (call.open) {
const code = /^RESTRICTED_[A-Z_]+$/.test(error.code ?? "") ? error.code : "RESTRICTED_PACKAGE_FAILED";
@@ -412,7 +502,14 @@ export async function serveRestrictedRegistry(config, io = process) {
throw fixedError("INVALID_RESTRICTED_BUDGET");
if (active.has(message.invocation) || queued.some((call) => call.message.invocation === message.invocation))
throw fixedError("RESTRICTED_DUPLICATE_INVOCATION");
const call = { message, open: true, pending: new Map(), sequence: 0, nested: (message.budget?.depth ?? 0) > 1 };
const call = {
message,
receivedAt: performance.now(),
open: true,
pending: new Map(),
sequence: 0,
nested: (message.budget?.depth ?? 0) > 1,
};
// Descendants cannot wait behind ancestors that hold all nested capacity.
// Reject excess fan-out/depth rather than deadlock an admitted call tree.
if (!available(call) && (call.nested || queued.length >= limits.maxQueued)) {
@@ -494,7 +591,16 @@ export async function serveRestrictedRegistry(config, io = process) {
io.stdin.on("error", (error) => void fail(error));
io.stdout.on("error", (error) => void fail(error));
send({ kind: "ready", exports: config.exports });
return { close: () => fail(fixedError("SERVICE_SHUTDOWN")) };
return {
close: () => fail(fixedError("SERVICE_SHUTDOWN")),
stats: () => ({
...metrics,
active: active.size,
queued: queued.length,
engines: engines.size,
idle: idleEngines.length,
}),
};
}
if (process.argv[1] === new URL(import.meta.url).pathname) {