diff --git a/server/src/services/remote-workspace-archive-worker.test.ts b/server/src/services/remote-workspace-archive-worker.test.ts new file mode 100644 index 00000000000..e608be9aedd --- /dev/null +++ b/server/src/services/remote-workspace-archive-worker.test.ts @@ -0,0 +1,107 @@ +import { createHash } from "node:crypto"; +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const state = vi.hoisted(() => ({ control: undefined as SharedArrayBuffer | undefined, + mode: "" as "" | "cpu" | "abort", spawned: 0, exited: 0 })); +vi.mock("node:worker_threads", async importOriginal => { + const actual = await importOriginal(); + return { ...actual, Worker: class extends actual.Worker { + constructor(source: string | URL, options?: import("node:worker_threads").WorkerOptions) { + if (state.control && typeof source === "string") { + const setup = `const testControl=new Int32Array(workerData.testControl);`; + source = source.replace("const archive=", `${setup}const archive=`); + if (state.mode === "cpu") source = source.replace("const before = await walk(root, true);", ` + Atomics.add(testControl,0,1); + while(!Atomics.load(testControl,1)) Atomics.add(testControl,2,1); + const before = await walk(root, true);`); + if (state.mode === "abort") { + source = source.replace("const before = await walk(root, true);", ` + Atomics.store(testControl,0,1); + await new Promise(resolve=>{input.signal.addEventListener('abort',resolve,{once:true});if(input.signal.aborted)resolve();}); + check(); + const before = await walk(root, true);`); + source = source.replace(/await fs\.rm\(scratch,\s*\{\s*recursive:\s*true,\s*force:\s*true\s*\}\);?/, ` + Atomics.store(testControl,3,1); + while(!Atomics.load(testControl,4)) await new Promise(resolve=>setTimeout(resolve,5)); + await fs.rm(scratch, { recursive: true, force: true });`); + } + options = { ...options, workerData: { ...options?.workerData, testControl: state.control } }; + } + super(source, options); + state.spawned++; + this.once("exit", () => { state.exited++; }); + } + } }; +}); +import { captureRemoteWorkspaceDirectory, restoreRemoteWorkspaceArchive, workspaceArchive } from "./remote-workspace-snapshot-archive.js"; + +describe("workspace archive worker boundary", () => { + let root: string, cwd: string; + beforeEach(async () => { + state.control = undefined; state.mode = ""; state.spawned = 0; state.exited = 0; + root = await fs.mkdtemp(path.join(os.tmpdir(), "pc-archive-worker-test-")); + cwd = path.join(root, "source"); await fs.mkdir(cwd); + await fs.writeFile(path.join(cwd, "saved.txt"), "retained work\n"); + await fs.mkdir(path.join(cwd, "nested")); + await fs.writeFile(path.join(cwd, "nested", "bytes"), Buffer.alloc(1024 * 1024, 42)); + await fs.chmod(path.join(cwd, "saved.txt"), 0o751); + await fs.symlink("saved.txt", path.join(cwd, "link")); + }); + afterEach(async () => { + if (state.control) { const control = new Int32Array(state.control); Atomics.store(control,1,1); Atomics.store(control,4,1); } + await fs.rm(root, { recursive: true, force: true }); + }); + it("preserves exact archive bytes, manifest and restored files across real workers", async () => { + const directPath = path.join(root, "direct.pcws"), archivePath = path.join(root, "worker.pcws"); + const direct = await workspaceArchive({ kind: "capture", cwd, archivePath: directPath }); + const captured = await captureRemoteWorkspaceDirectory({ cwd, archivePath }); + expect(captured).toEqual(direct); + expect(await fs.readFile(archivePath)).toEqual(await fs.readFile(directPath)); + const targetDir = path.join(root, "restored"); + await expect(restoreRemoteWorkspaceArchive({ archivePath, targetDir, expectedManifestDigest: direct.manifestDigest })).resolves.toEqual(direct); + expect(await fs.readFile(path.join(targetDir, "saved.txt"), "utf8")).toBe("retained work\n"); + expect((await fs.stat(path.join(targetDir, "saved.txt"))).mode & 0o777).toBe(0o751); + expect(await fs.readlink(path.join(targetDir, "link"))).toBe("saved.txt"); + expect(createHash("sha256").update(await fs.readFile(path.join(targetDir, "nested", "bytes"))).digest("hex")) + .toBe(createHash("sha256").update(await fs.readFile(path.join(cwd, "nested", "bytes"))).digest("hex")); + expect(state.spawned).toBe(2); expect(state.exited).toBe(2); + }); + it("runs CPU barriers on eight real workers while an aborted queued capture never spawns", async () => { + state.control = new SharedArrayBuffer(5 * 4); state.mode = "cpu"; + const control = new Int32Array(state.control), controller = new AbortController(), reason = new Error("queued capture cancelled"); + const captures = Array.from({ length: 8 }, (_, i) => captureRemoteWorkspaceDirectory({ cwd, archivePath: path.join(root, `${i}.pcws`) })); + const queued = captureRemoteWorkspaceDirectory({ cwd, archivePath: path.join(root, "queued.pcws"), signal: controller.signal }).catch(error => error); + try { + await expect.poll(() => Atomics.load(control,0), { timeout: 10_000 }).toBe(8); + expect(state.spawned).toBe(8); expect(state.exited).toBe(0); + expect(Atomics.load(control,2)).toBeGreaterThan(0); + controller.abort(reason); + expect(await queued).toBe(reason); expect(state.spawned).toBe(8); + Atomics.store(control,1,1); + await Promise.all(captures); + expect(state.exited).toBe(8); + await expect(fs.access(path.join(root, "queued.pcws"))).rejects.toMatchObject({ code: "ENOENT" }); + } finally { Atomics.store(control,1,1); await Promise.allSettled(captures); } + }); + it("waits for abort cleanup and actual worker exit before rejecting or releasing its slot", async () => { + state.control = new SharedArrayBuffer(5 * 4); state.mode = "abort"; + const control = new Int32Array(state.control), controller = new AbortController(), reason = new Error("owned capture cancelled"); + let settled = false; + const archivePath = path.join(root, "cancelled.pcws"); + const capture = captureRemoteWorkspaceDirectory({ cwd, archivePath, signal: controller.signal }).then( + () => { settled = true; return null; }, error => { settled = true; return error; }); + try { + await expect.poll(() => Atomics.load(control,0), { timeout: 10_000 }).toBe(1); + controller.abort(reason); + await expect.poll(() => Atomics.load(control,3), { timeout: 10_000 }).toBe(1); + expect(settled).toBe(false); expect(state.exited).toBe(0); + Atomics.store(control,4,1); + expect(await capture).toBe(reason); + expect(state.exited).toBe(1); + await expect(fs.access(archivePath)).rejects.toMatchObject({ code: "ENOENT" }); + } finally { Atomics.store(control,4,1); controller.abort(reason); await capture; } + }); +}); diff --git a/server/src/services/remote-workspace-snapshot-archive.test.ts b/server/src/services/remote-workspace-snapshot-archive.test.ts index 8b49be22069..3401d7e1aa7 100644 --- a/server/src/services/remote-workspace-snapshot-archive.test.ts +++ b/server/src/services/remote-workspace-snapshot-archive.test.ts @@ -6,7 +6,7 @@ import os from "node:os"; import path from "node:path"; import { promisify } from "node:util"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { captureRemoteWorkspaceDirectory, readRemoteWorkspaceArchiveDirectories, remoteWorkspaceArchiveScript, remoteWorkspaceCaptureFailureMessage, restoreRemoteWorkspaceArchive } from "./remote-workspace-snapshot-archive.js"; +import { workspaceArchive, readRemoteWorkspaceArchiveDirectories, remoteWorkspaceArchiveScript, remoteWorkspaceCaptureFailureMessage } from "./remote-workspace-snapshot-archive.js"; const archiveIo = vi.hoisted(() => ({ afterChmod: undefined as ((target: string, mode: number) => void) | undefined, @@ -35,6 +35,9 @@ vi.mock("node:fs/promises", async importOriginal => { } }; }); +const captureRemoteWorkspaceDirectory = (input: Omit[0], "kind">) => workspaceArchive({ ...input, kind: "capture" }); +const restoreRemoteWorkspaceArchive = (input: Omit[0], "kind">) => workspaceArchive({ ...input, kind: "restore" }); + const execute = promisify(execFile); const env = { ...process.env, GIT_CONFIG_NOSYSTEM: "1", GIT_CONFIG_GLOBAL: os.devNull, GIT_AUTHOR_NAME: "Snapshot", GIT_COMMITTER_NAME: "Snapshot", GIT_AUTHOR_EMAIL: "snapshot@example.test", GIT_COMMITTER_EMAIL: "snapshot@example.test" }; diff --git a/server/src/services/remote-workspace-snapshot-archive.ts b/server/src/services/remote-workspace-snapshot-archive.ts index 3bed12712b1..9af1441c804 100644 --- a/server/src/services/remote-workspace-snapshot-archive.ts +++ b/server/src/services/remote-workspace-snapshot-archive.ts @@ -1,8 +1,10 @@ +import { Worker } from "node:worker_threads"; + /** * A portable, streaming workspace archive. This function is intentionally * self-contained: the same implementation runs locally and in a sandbox. */ -async function workspaceArchive(input: { +export async function workspaceArchive(input: { kind: "capture" | "restore"; cwd?: string; archivePath: string; @@ -662,8 +664,72 @@ async function workspaceArchive(input: { } } +type ArchiveInput = Parameters[0]; +let activeArchiveWorkers = 0; +const archiveWaiters: Array<() => void> = []; + +async function archiveWorkerSlot(signal?: AbortSignal): Promise<() => void> { + signal?.throwIfAborted(); + await new Promise((resolve, reject) => { + const start = () => { + signal?.removeEventListener("abort", abort); + activeArchiveWorkers++; + resolve(); + }; + const abort = () => { + const index = archiveWaiters.indexOf(start); + if (index >= 0) archiveWaiters.splice(index, 1); + reject(signal!.reason); + }; + if (activeArchiveWorkers < 8) start(); + else { + archiveWaiters.push(start); + signal?.addEventListener("abort", abort, { once: true }); + } + }); + return () => { activeArchiveWorkers--; archiveWaiters.shift()?.(); }; +} + +async function runArchiveWorker(input: ArchiveInput): ReturnType { + const release = await archiveWorkerSlot(input.signal); + try { + input.signal?.throwIfAborted(); + const { signal, ...parameters } = input; + const source = workspaceArchive.toString().replaceAll("__vite_ssr_dynamic_import__", "import"); + const worker = new Worker(`const {parentPort,workerData}=require("node:worker_threads"); +const __name=value=>value; +const controller=new AbortController(); +parentPort.on("message",()=>controller.abort()); +const archive=${source}; +archive({...workerData,signal:controller.signal}).then(result=>parentPort.postMessage({result}),error=>parentPort.postMessage({error:{...error,name:error.name,message:error.message}})).finally(()=>parentPort.close());`, { + eval: true, workerData: parameters, execArgv: [], stdout: true, stderr: true, + resourceLimits: { maxOldGenerationSizeMb: 1024, maxYoungGenerationSizeMb: 64, stackSizeMb: 4 }, + }); + worker.stdout.resume(); worker.stderr.resume(); + const abort = () => worker.postMessage("abort"); + signal?.addEventListener("abort", abort, { once: true }); + if (signal?.aborted) abort(); + try { + return await new Promise>>((resolve, reject) => { + let result: Awaited> | undefined, failure: Error | undefined; + worker.once("message", (message: { result?: typeof result; error?: Record }) => { + result = message.result; + if (message.error) failure = Object.assign(new Error(String(message.error.message)), message.error); + }); + worker.once("error", error => { failure = error; }); + worker.once("exit", code => { + if (signal?.aborted) reject(signal.reason); + else if (failure) reject(failure); + else if (code === 0 && result) resolve(result); + else reject(new Error("remote_workspace_archive_worker_failed")); + }); + }); + } finally { signal?.removeEventListener("abort", abort); } + } finally { release(); } +} + export function captureRemoteWorkspaceDirectory(input: { cwd: string; archivePath: string; manifestPath?: string; exclude?: string[]; requireNoGit?: boolean; gitSourceRoots?: string[]; signal?: AbortSignal }) { - return workspaceArchive({ ...input, kind: "capture" }); + return runArchiveWorker({ ...input, kind: "capture" }); } export async function readRemoteWorkspaceArchiveDirectories(archivePath: string, expectedManifestDigest: string) { @@ -691,7 +757,7 @@ export async function readRemoteWorkspaceArchiveDirectories(archivePath: string, } export function restoreRemoteWorkspaceArchive(input: { archivePath: string; targetDir: string; expectedManifestDigest: string; signal?: AbortSignal; maxEntries?: number; maxManifestBytes?: number }) { - return workspaceArchive({ ...input, kind: "restore" }); + return runArchiveWorker({ ...input, kind: "restore" }); } export function remoteWorkspaceCaptureFailureMessage(error: unknown): string | null {