@ricsam/r5d-worker 0.0.109 → 0.0.111

package/dist/mjs/job-resources.mjs

@@ -0,0 +1,293 @@
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
+const MIB = 1024 * 1024;
+const WORKER_INLINE_OUTPUT_MAX_BYTES = 8 * MIB;
+async function collectWorkerInlineOutput(stream, maxBytes = WORKER_INLINE_OUTPUT_MAX_BYTES, signal) {
+  if (!stream) return "";
+  signal?.throwIfAborted();
+  const reader = stream.getReader();
+  const abort = () => {
+    void reader.cancel().catch(() => {
+    });
+  };
+  signal?.addEventListener("abort", abort, { once: true });
+  const decoder = new TextDecoder();
+  const parts = [];
+  let bytes = 0;
+  try {
+    while (true) {
+      const next = await reader.read();
+      signal?.throwIfAborted();
+      if (next.done) break;
+      bytes += next.value.byteLength;
+      if (bytes > maxBytes) {
+        void reader.cancel().catch(() => {
+        });
+        throw new Error(`Inline command output exceeded ${maxBytes} bytes. Use streaming execution for large output.`);
+      }
+      parts.push(decoder.decode(next.value, { stream: true }));
+    }
+    parts.push(decoder.decode());
+    return parts.join("");
+  } finally {
+    signal?.removeEventListener("abort", abort);
+    reader.releaseLock();
+  }
+}
+function positiveInteger(value, fallback, name) {
+  const result = value === void 0 ? fallback : Number(value);
+  if (!Number.isSafeInteger(result) || result <= 0) throw new Error(`${name} must be a positive integer`);
+  return result;
+}
+function workerJobBudget(effectiveMemoryBytes, env = process.env) {
+  const generalSlots = positiveInteger(env.R5D_WORKER_JOB_CONCURRENCY, 2, "R5D_WORKER_JOB_CONCURRENCY");
+  const utilitySlots = positiveInteger(env.R5D_WORKER_UTILITY_CONCURRENCY, 2, "R5D_WORKER_UTILITY_CONCURRENCY");
+  const memoryPercent = positiveInteger(env.R5D_WORKER_JOBS_MEMORY_PERCENT, 70, "R5D_WORKER_JOBS_MEMORY_PERCENT");
+  if (memoryPercent > 70) throw new Error("Worker jobs may reserve at most 70% of effective memory");
+  const aggregateBytes = positiveInteger(
+    env.R5D_WORKER_JOBS_MEMORY_BYTES,
+    Math.floor(effectiveMemoryBytes * memoryPercent / 100),
+    "R5D_WORKER_JOBS_MEMORY_BYTES"
+  );
+  const utilityBytes = positiveInteger(
+    env.R5D_WORKER_UTILITY_MEMORY_BYTES,
+    Math.min(256 * MIB, Math.floor(aggregateBytes / 10)),
+    "R5D_WORKER_UTILITY_MEMORY_BYTES"
+  );
+  const generalBytes = positiveInteger(
+    env.R5D_WORKER_JOB_MEMORY_BYTES,
+    Math.floor((aggregateBytes - utilityBytes * utilitySlots) / generalSlots),
+    "R5D_WORKER_JOB_MEMORY_BYTES"
+  );
+  if (aggregateBytes > Math.floor(effectiveMemoryBytes * 0.7)) throw new Error("Worker jobs may reserve at most 70% of effective memory");
+  if (generalBytes * generalSlots + utilityBytes * utilitySlots > aggregateBytes)
+    throw new Error("Worker job limits exceed the aggregate memory budget");
+  const queueTimeoutMs = positiveInteger(env.R5D_WORKER_JOB_QUEUE_TIMEOUT_MS, 6e4, "R5D_WORKER_JOB_QUEUE_TIMEOUT_MS");
+  if (queueTimeoutMs > 6e4) throw new Error("Worker job queue timeout cannot exceed 60000ms");
+  return {
+    aggregateBytes,
+    generalBytes,
+    utilityBytes,
+    generalSlots,
+    utilitySlots,
+    queueTimeoutMs,
+    maxQueued: positiveInteger(env.R5D_WORKER_JOB_QUEUE_LIMIT, 128, "R5D_WORKER_JOB_QUEUE_LIMIT"),
+    pidsPerJob: positiveInteger(env.R5D_WORKER_JOB_PIDS, 512, "R5D_WORKER_JOB_PIDS")
+  };
+}
+function workerDelegatedCgroupRoot(cgroupText, mount = "/sys/fs/cgroup") {
+  const relative = cgroupText.split("\n").find((line) => line.startsWith("0::/"))?.slice(3);
+  if (!relative) return void 0;
+  const segments = relative.split("/").filter(Boolean);
+  const index = segments.indexOf("r5d-worker.service");
+  if (index < 0 || segments[index + 1] !== "control") return void 0;
+  return path.join(mount, ...segments.slice(0, index + 1));
+}
+function effectiveWorkerMemory(root, hostBytes = os.totalmem(), mount = "/sys/fs/cgroup") {
+  let result = hostBytes;
+  if (!root) return result;
+  let current = root;
+  while (current === mount || current.startsWith(`${mount}/`)) {
+    try {
+      const limit = Number(fs.readFileSync(path.join(current, "memory.max"), "utf8").trim());
+      if (Number.isSafeInteger(limit) && limit > 0) result = Math.min(result, limit);
+    } catch {
+    }
+    if (current === mount) break;
+    current = path.dirname(current);
+  }
+  return result;
+}
+class WorkerJobCapacityError extends Error {
+  constructor() {
+    super("Worker command capacity is busy. No command was started; retry later or stop another command or terminal.");
+    this.name = "WorkerJobCapacityError";
+  }
+}
+class WorkerJobCancelledError extends Error {
+  constructor() {
+    super("Command cancelled before it started");
+    this.name = "WorkerJobCancelledError";
+  }
+}
+class WorkerJobResourceController {
+  constructor(budget, delegatedRoot) {
+    this.budget = budget;
+    this.isolation = delegatedRoot ? "cgroup-v2" : "concurrency-only";
+    if (delegatedRoot) {
+      if (fs.readFileSync(path.join(delegatedRoot, "cgroup.procs"), "utf8").trim())
+        throw new Error("Worker cgroup delegation root must be empty; configure DelegateSubgroup=control");
+      fs.writeFileSync(path.join(delegatedRoot, "cgroup.subtree_control"), "+memory +pids");
+      this.jobsRoot = path.join(delegatedRoot, "jobs");
+      fs.mkdirSync(this.jobsRoot, { recursive: true });
+      fs.writeFileSync(path.join(this.jobsRoot, "memory.max"), String(budget.aggregateBytes));
+      fs.writeFileSync(path.join(this.jobsRoot, "memory.swap.max"), "0");
+      fs.writeFileSync(path.join(this.jobsRoot, "memory.oom.group"), "0");
+      fs.writeFileSync(path.join(this.jobsRoot, "cgroup.subtree_control"), "+memory +pids");
+      for (const name of fs.readdirSync(this.jobsRoot)) {
+        if (!/^job-[0-9a-f-]+$/.test(name)) continue;
+        const old = path.join(this.jobsRoot, name);
+        if (!fs.statSync(old).isDirectory()) continue;
+        fs.writeFileSync(path.join(old, "cgroup.kill"), "1");
+      }
+    }
+  }
+  budget;
+  isolation;
+  jobsRoot;
+  active = /* @__PURE__ */ new Map();
+  waiters = [];
+  timer;
+  get status() {
+    return { isolation: this.isolation, active: this.active.size, queued: this.waiters.length };
+  }
+  acquire(input) {
+    try {
+      input.assertAdmission();
+      if (this.active.has(input.id) || this.waiters.some((entry) => entry.id === input.id))
+        throw new Error(`Duplicate resource admission: ${input.id}`);
+      if (this.available(input.lane)) return Promise.resolve(this.start(input));
+      if (this.waiters.length >= this.budget.maxQueued) throw new WorkerJobCapacityError();
+      return new Promise((resolve, reject) => {
+        this.waiters.push({ ...input, resolve, reject, deadline: Date.now() + this.budget.queueTimeoutMs });
+        this.timer ??= setInterval(() => this.drain(), Math.min(100, this.budget.queueTimeoutMs));
+      });
+    } catch (error) {
+      return Promise.reject(error);
+    }
+  }
+  cancel(id) {
+    for (let index = this.waiters.length - 1; index >= 0; index--) {
+      if (this.waiters[index].id !== id) continue;
+      this.waiters.splice(index, 1)[0].reject(new WorkerJobCancelledError());
+    }
+    this.active.get(id)?.lease.kill();
+    this.drain();
+  }
+  cancelSession(sessionId) {
+    this.cancelMatching((_id, ownerSessionId) => ownerSessionId === sessionId);
+  }
+  cancelMatching(matches) {
+    const ids = /* @__PURE__ */ new Set([
+      ...this.waiters.filter((entry) => matches(entry.id, entry.sessionId)).map((entry) => entry.id),
+      ...[...this.active].filter(([id, entry]) => matches(id, entry.sessionId)).map(([id]) => id)
+    ]);
+    const errors = [];
+    for (const id of ids) {
+      try {
+        this.cancel(id);
+      } catch (error) {
+        errors.push(error);
+      }
+    }
+    if (errors.length) throw new AggregateError(errors, "Worker job cancellation was incomplete");
+  }
+  available(lane) {
+    return [...this.active.values()].filter((entry) => entry.lane === lane).length < (lane === "general" ? this.budget.generalSlots : this.budget.utilitySlots);
+  }
+  drain() {
+    for (let index = 0; index < this.waiters.length; ) {
+      const waiter = this.waiters[index];
+      try {
+        waiter.assertAdmission();
+        if (Date.now() >= waiter.deadline) throw new WorkerJobCapacityError();
+        if (!this.available(waiter.lane)) {
+          index++;
+          continue;
+        }
+        this.waiters.splice(index, 1);
+        waiter.resolve(this.start(waiter));
+      } catch (error) {
+        if (this.waiters[index] === waiter) this.waiters.splice(index, 1);
+        waiter.reject(error);
+      }
+    }
+    if (this.waiters.length === 0 && this.timer) {
+      clearInterval(this.timer);
+      this.timer = void 0;
+    }
+  }
+  start(input) {
+    input.assertAdmission();
+    const memoryBytes = input.lane === "general" ? this.budget.generalBytes : this.budget.utilityBytes;
+    const cgroupPath = this.jobsRoot ? path.join(this.jobsRoot, `job-${crypto.randomUUID()}`) : void 0;
+    if (cgroupPath) {
+      fs.mkdirSync(cgroupPath);
+      try {
+        fs.writeFileSync(path.join(cgroupPath, "memory.max"), String(memoryBytes));
+        fs.writeFileSync(path.join(cgroupPath, "memory.swap.max"), "0");
+        fs.writeFileSync(path.join(cgroupPath, "memory.oom.group"), "1");
+        fs.writeFileSync(path.join(cgroupPath, "pids.max"), String(this.budget.pidsPerJob));
+      } catch (error) {
+        fs.rmdirSync(cgroupPath);
+        throw error;
+      }
+    }
+    let released = false;
+    let revoked = false;
+    let releasePromise;
+    const kill = () => {
+      revoked = true;
+      if (!released && cgroupPath) fs.writeFileSync(path.join(cgroupPath, "cgroup.kill"), "1");
+    };
+    const lease = {
+      cgroupPath,
+      memoryBytes,
+      wrap: (argv) => {
+        if (released) throw new Error("Worker job resource lease was already released");
+        if (revoked) throw new WorkerJobCancelledError();
+        if (!argv.length) throw new Error("Command argv must not be empty");
+        return cgroupPath ? [
+          "/bin/sh",
+          "-c",
+          'set -e; printf "%s" "$$" > "$1/cgroup.procs"; read -r oom < /proc/self/oom_score_adj; if [ "$oom" -lt 0 ]; then printf 0 > /proc/self/oom_score_adj; fi; shift; exec "$@"',
+          "r5d-job",
+          cgroupPath,
+          ...argv
+        ] : argv;
+      },
+      kill,
+      memoryLimitExceeded: () => {
+        if (!cgroupPath || released) return false;
+        return /^oom_kill [1-9][0-9]*$/m.test(fs.readFileSync(path.join(cgroupPath, "memory.events"), "utf8"));
+      },
+      release: () => releasePromise ??= (async () => {
+        kill();
+        if (cgroupPath) {
+          const deadline = Date.now() + 2e3;
+          while (/^populated 1$/m.test(fs.readFileSync(path.join(cgroupPath, "cgroup.events"), "utf8"))) {
+            if (Date.now() >= deadline) throw new Error(`Worker job cgroup remained populated: ${cgroupPath}`);
+            await new Promise((resolve) => setTimeout(resolve, 10));
+          }
+          fs.rmdirSync(cgroupPath);
+        }
+        released = true;
+        this.active.delete(input.id);
+        this.drain();
+      })()
+    };
+    this.active.set(input.id, { lane: input.lane, sessionId: input.sessionId, lease });
+    return lease;
+  }
+}
+function createWorkerJobResources(env = process.env) {
+  const mode = env.R5D_WORKER_RESOURCE_ISOLATION ?? "auto";
+  if (mode !== "auto" && mode !== "required") throw new Error("R5D_WORKER_RESOURCE_ISOLATION must be auto or required");
+  const root = process.platform === "linux" ? workerDelegatedCgroupRoot(fs.readFileSync("/proc/self/cgroup", "utf8")) : void 0;
+  if (mode === "required" && !root)
+    throw new Error("Worker requires delegated cgroup v2 resource isolation; configure Delegate=memory pids and DelegateSubgroup=control");
+  const budget = workerJobBudget(effectiveWorkerMemory(root), env);
+  return new WorkerJobResourceController(budget, root);
+}
+export {
+  WORKER_INLINE_OUTPUT_MAX_BYTES,
+  WorkerJobCancelledError,
+  WorkerJobCapacityError,
+  WorkerJobResourceController,
+  collectWorkerInlineOutput,
+  createWorkerJobResources,
+  effectiveWorkerMemory,
+  workerDelegatedCgroupRoot,
+  workerJobBudget
+};