fix(artifacts): isolate cancelled PyInstaller parse attempts (#9961)

src/artifacts/PyInstallerArtifactReader.ts

@@ -23,6 +23,7 @@ import {
   type ArtifactTraversalOptions,
 } from "./ArtifactReader.js";
 import { abortIfNeeded } from "./ArtifactHash.js";
+import { artifactReadCancelledBy } from "./ArtifactFileOperation.js";
 import { ArtifactFileSource } from "./ArtifactFileSource.js";
 import { ArtifactResourceScope } from "./ArtifactResourceScope.js";
 import { parseMzHeaderEvidence } from "../domain/dosMz.js";
@@ -99,7 +100,17 @@ const carrierPathDialect = async (
   return "unknown";
 };
 
+interface ArchiveParseState {
+  readonly signal: AbortSignal | undefined;
+  readonly pyzStore: PyInstallerPyzStore;
+  budget: number;
+  charged: number;
+  parsedFactBytes: number;
+  retainedEntryBytes: number;
+}
+
 interface LoadedArchive {
+  readonly state: ArchiveParseState;
   readonly members: readonly PyInstallerMember[];
   readonly traversal: ArtifactReaderTraversal;
 }
@@ -116,13 +127,13 @@ export class PyInstallerArtifactReader implements ArtifactReader {
   readonly #scope = new ArtifactResourceScope();
   readonly #cleanup = new ArtifactResourceScope();
   readonly #lifecycle = new AbortController();
-  #loaded: Promise<LoadedArchive> | undefined;
-  #budget = 0;
-  #charged = 0;
-  #parsedFactBytes = 0;
-  #retainedEntryBytes = 0;
+  #loaded:
+    | {
+        readonly signal: AbortSignal | undefined;
+        readonly promise: Promise<LoadedArchive>;
+      }
+    | undefined;
   #metadata: PyInstallerMetadata | undefined;
-  readonly #pyzStore = new PyInstallerPyzStore();
   readonly #members = new Map<string, PyInstallerMember>();
   readonly #openedPyzSizes = new Map<string, number>();
 
@@ -136,11 +147,6 @@ export class PyInstallerArtifactReader implements ArtifactReader {
     } else {
       this.#source = new ArtifactFileSource(path, "PyInstaller", source);
     }
-    this.#cleanup.retain({
-      kind: "lease",
-      resource: "PyInstaller private PYZ images",
-      lease: { release: () => this.#pyzStore.close() },
-    });
     this.#cleanup.retain({
       kind: "lease",
       resource: path,
@@ -202,6 +208,7 @@ export class PyInstallerArtifactReader implements ArtifactReader {
           "unavailable",
           `PYZ entry is encrypted and was not decoded: ${entry.path}`,
         );
+      const { state } = await this.#load(active);
       const read = stored.read;
       if (read === undefined)
         throw new ArtifactReaderFailure(
@@ -215,8 +222,8 @@ export class PyInstallerArtifactReader implements ArtifactReader {
             this.#openedPyzSizes.get(entry.adapterKey))
           : undefined;
       if (read.kind === "pyz" && knownSize === undefined)
-        this.#charge(read.header?.length ?? 0);
-      const parts = await this.#pyzStore.memberSource(read, handle, active);
+        this.#charge(state, read.header?.length ?? 0);
+      const parts = await state.pyzStore.memberSource(read, handle, active);
       if (active.aborted) {
         parts.source.destroy();
         await finished(parts.source, { cleanup: true }).catch(() => undefined);
@@ -240,10 +247,10 @@ export class PyInstallerArtifactReader implements ArtifactReader {
         maxBytes:
           read.kind === "carchive"
             ? read.source.uncompressedSize
-            : (knownSize ?? this.#budget - this.#charged),
+            : (knownSize ?? state.budget - state.charged),
         signal: active,
         ...(read.kind === "pyz" && knownSize === undefined
-          ? { charge: (bytes: number) => this.#charge(bytes) }
+          ? { charge: (bytes: number) => this.#charge(state, bytes) }
           : {}),
         complete: async (bytes) => {
           try {
@@ -315,11 +322,60 @@ export class PyInstallerArtifactReader implements ArtifactReader {
   }
 
   async #load(signal?: AbortSignal): Promise<LoadedArchive> {
-    this.#loaded ??= this.#parse(signal);
-    return this.#loaded;
+    for (;;) {
+      abortIfNeeded(signal);
+      const pending = (this.#loaded ??= {
+        signal,
+        promise: this.#parseAttempt(signal),
+      });
+      let loaded: LoadedArchive;
+      try {
+        loaded = await pending.promise;
+      } catch (cause: unknown) {
+        if (artifactReadCancelledBy(cause, pending.signal)) {
+          if (this.#loaded === pending) this.#loaded = undefined;
+          if (pending.signal !== signal && signal?.aborted !== true) continue;
+        }
+        throw cause;
+      }
+      abortIfNeeded(signal);
+      return loaded;
+    }
   }
 
-  async #parse(signal?: AbortSignal): Promise<LoadedArchive> {
+  async #parseAttempt(signal?: AbortSignal): Promise<LoadedArchive> {
+    const state: ArchiveParseState = {
+      signal,
+      pyzStore: new PyInstallerPyzStore(),
+      budget: 0,
+      charged: 0,
+      parsedFactBytes: 0,
+      retainedEntryBytes: 0,
+    };
+    const owner = {
+      kind: "lease" as const,
+      resource: "PyInstaller private PYZ images",
+      lease: { release: () => state.pyzStore.close() },
+    };
+    this.#cleanup.retain(owner);
+    try {
+      return await this.#parse(state);
+    } catch (cause: unknown) {
+      // A cancelled attempt owns its private files until their cleanup settles.
+      // Failed cleanup stays with the reader and blocks reuse of the attempt.
+      if (artifactReadCancelledBy(cause, signal)) {
+        const cleanup = await this.#cleanup.release(owner);
+        if (cleanup.kind === "failed")
+          throw this.#cleanup.cleanupFailure({ kind: "failed", cause }, [
+            cleanup,
+          ]);
+      }
+      throw cause;
+    }
+  }
+
+  async #parse(state: ArchiveParseState): Promise<LoadedArchive> {
+    const { signal } = state;
     await this.#openHandle(signal);
     const { handle, initial } = await this.#source.readSelectedPath();
     const size = initial.size;
@@ -328,37 +384,39 @@ export class PyInstallerArtifactReader implements ArtifactReader {
       (await carrierPathDialect(handle, size, signal));
     const cookie =
       this.#admission?.cookie ?? (await this.#cookie(handle, size, signal));
-    this.#budget = artifactExpansionBudgetBytes(cookie.packageLength);
+    state.budget = artifactExpansionBudgetBytes(cookie.packageLength);
     const toc = await this.#toc(handle, cookie, signal);
     for (const entry of toc) {
-      this.#charge(entry.uncompressedSize);
-      this.#reserveMetadata({ ...entry }, false);
+      this.#charge(state, entry.uncompressedSize);
+      this.#reserveMetadata(state, { ...entry }, false);
     }
-    const images = await this.#pyzImages(handle, cookie, toc, signal);
+    const images = await this.#pyzImages(state, handle, cookie, toc);
     const magic = scriptPythonMagic(
       cookie,
       [...images.values()].map((image) => image.archive),
     );
     const encrypted = hasEncryptionMarker(toc);
     const members = await this.#project(
+      state,
       toc,
       cookie,
       images,
       magic,
       encrypted,
       pathDialect,
-      signal,
     );
     for (const member of members)
       if (member.read?.kind === "carchive")
-        this.#charge(member.read.header?.length ?? 0);
+        this.#charge(state, member.read.header?.length ?? 0);
     await this.#source.readSelectedPath();
+    abortIfNeeded(signal);
     this.#admission = undefined;
     this.#metadata = cookieMetadata(cookie, toc, encrypted);
     let opaqueImages = 0;
     for (const image of images.values())
       if (image.archive.kind === "opaque") opaqueImages++;
     return {
+      state,
       members,
       traversal: {
         status: opaqueImages === 0 ? "complete" : "unknown",
@@ -430,14 +488,15 @@ export class PyInstallerArtifactReader implements ArtifactReader {
   }
 
   async #pyzImages(
+    state: ArchiveParseState,
     handle: FileHandle,
     cookie: PyInstallerCookie,
     toc: readonly CArchiveEntry[],
-    signal: AbortSignal | undefined,
   ): Promise<ReadonlyMap<CArchiveEntry, PyzImage>> {
+    const { signal } = state;
     const images = new Map<CArchiveEntry, PyzImage>();
     // Admit the complete filesystem footprint before the first private write.
-    this.#pyzStore.reserveImages(
+    state.pyzStore.reserveImages(
       toc.reduce(
         (count, entry) =>
           count + (entry.role === "pyz" && entry.compressionFlag === 1 ? 1 : 0),
@@ -450,7 +509,7 @@ export class PyInstallerArtifactReader implements ArtifactReader {
       const spooled =
         entry.compressionFlag === 0
           ? undefined
-          : await this.#pyzStore.spool(
+          : await state.pyzStore.spool(
               handle,
               cookie.archiveStart,
               entry,
@@ -469,8 +528,12 @@ export class PyInstallerArtifactReader implements ArtifactReader {
       if (!parsed.ok) throw failure(parsed.error);
       if (parsed.value.kind === "entries")
         for (const item of parsed.value.entries)
-          this.#reserveMetadata({ ...item, container: entry.name }, false);
-      if (spooled !== undefined) await this.#pyzStore.release(spooled.handle);
+          this.#reserveMetadata(
+            state,
+            { ...item, container: entry.name },
+            false,
+          );
+      if (spooled !== undefined) await state.pyzStore.release(spooled.handle);
       images.set(entry, {
         archive: parsed.value,
         handle: source.handle,
@@ -481,23 +544,25 @@ export class PyInstallerArtifactReader implements ArtifactReader {
   }
 
   async #project(
+    state: ArchiveParseState,
     toc: readonly CArchiveEntry[],
     cookie: PyInstallerCookie,
     images: ReadonlyMap<CArchiveEntry, PyzImage>,
     magic: Buffer | undefined,
     encrypted: boolean,
     pathDialect: CarrierPathDialect,
-    signal: AbortSignal | undefined,
   ): Promise<PyInstallerMember[]> {
+    const { signal } = state;
     const handle = await this.#openHandle(signal);
     const members: PyInstallerMember[] = [];
     const append = (member: PyInstallerMember): void => {
-      this.#reservePathComponents(member.entry.path);
+      this.#reservePathComponents(state, member.entry.path);
       const entry = {
         ...member.entry,
         path: normalizeArtifactPath(member.entry.path),
       };
       this.#reserveMetadata(
+        state,
         {
           path: entry.path,
           source_path: entry.sourcePath ?? null,
@@ -514,7 +579,7 @@ export class PyInstallerArtifactReader implements ArtifactReader {
       abortIfNeeded(signal);
       const image = entry.role === "pyz" ? images.get(entry) : undefined;
       if (image === undefined) {
-        const bytecode = await this.#pyzStore.bytecodeHeader(
+        const bytecode = await state.pyzStore.bytecodeHeader(
           handle,
           cookie.archiveStart,
           entry,
@@ -565,8 +630,8 @@ export class PyInstallerArtifactReader implements ArtifactReader {
     return members;
   }
 
-  #reservePathComponents(path: string): void {
-    const remaining = MAX_RETAINED_ENTRY_BYTES - this.#retainedEntryBytes;
+  #reservePathComponents(state: ArchiveParseState, path: string): void {
+    const remaining = MAX_RETAINED_ENTRY_BYTES - state.retainedEntryBytes;
     // Case annotations can carry the path and two prefixes at every component.
     // Reserve escaped text growth as well as registry/case trie nodes and maps.
     const componentCost =
@@ -584,11 +649,15 @@ export class PyInstallerArtifactReader implements ArtifactReader {
         "PyInstaller paths exceed their retained namespace budget",
         { boundary: "memory" },
       );
-    this.#retainedEntryBytes += reservation;
+    state.retainedEntryBytes += reservation;
   }
 
-  #reserveMetadata(value: JsonValue, projected: boolean): void {
-    const used = projected ? this.#retainedEntryBytes : this.#parsedFactBytes;
+  #reserveMetadata(
+    state: ArchiveParseState,
+    value: JsonValue,
+    projected: boolean,
+  ): void {
+    const used = projected ? state.retainedEntryBytes : state.parsedFactBytes;
     const limit = projected ? MAX_RETAINED_ENTRY_BYTES : MAX_PARSED_FACT_BYTES;
     const remaining = limit - used - RECORD_OWNERSHIP_RESERVATION;
     const bytes = remaining < 0 ? 0 : jsonValueByteLength(value, remaining);
@@ -599,21 +668,21 @@ export class PyInstallerArtifactReader implements ArtifactReader {
         { boundary: "memory" },
       );
     if (projected)
-      this.#retainedEntryBytes += bytes + RECORD_OWNERSHIP_RESERVATION;
-    else this.#parsedFactBytes += bytes + RECORD_OWNERSHIP_RESERVATION;
+      state.retainedEntryBytes += bytes + RECORD_OWNERSHIP_RESERVATION;
+    else state.parsedFactBytes += bytes + RECORD_OWNERSHIP_RESERVATION;
   }
 
-  #charge(bytes: number): void {
+  #charge(state: ArchiveParseState, bytes: number): void {
     if (
       !Number.isSafeInteger(bytes) ||
       bytes < 0 ||
-      this.#charged + bytes > this.#budget
+      state.charged + bytes > state.budget
     )
       throw new ArtifactReaderFailure(
         "limit",
-        `PyInstaller members declare ${String(this.#charged + Math.max(0, bytes))} uncompressed bytes, exceeding the decompression budget`,
+        `PyInstaller members declare ${String(state.charged + Math.max(0, bytes))} uncompressed bytes, exceeding the decompression budget`,
       );
-    this.#charged += bytes;
+    state.charged += bytes;
   }
 }
 

src/artifacts/pyinstaller/PyInstallerBoundaries.test.ts

@@ -13,6 +13,7 @@ import { hashReadable } from "../ArtifactHash.js";
 import { translateArtifactFailure } from "../ArtifactOperationFailure.js";
 import { projectAnalysisError } from "../../domain/analysisErrorProjection.js";
 import { marshalCode } from "../../python/PythonPyc.fixture.js";
+import * as artifactPaths from "../ArtifactPaths.js";
 import { PyInstallerPyzStore } from "./PyInstallerPyzStore.js";
 import {
   ArtifactReaderFailure,
@@ -861,3 +862,247 @@ it.each([false, true])(
     );
   },
 );
+
+it.each(
+  (["carrier", "private image"] as const).flatMap((stage) =>
+    ([false, true] as const).map((peer) => ({ stage, peer })),
+  ),
+)(
+  "isolates cancelled $stage parsing from a fresh inventory (concurrent: $peer)",
+  async ({ stage, peer }) => {
+    const payload = Buffer.from("body");
+    const bytes =
+      stage === "carrier"
+        ? archive(payload)
+        : cArchiveExecutable({
+            layout: "2.1",
+            pythonVersion: 309,
+            members: [
+              {
+                name: "PYZ.pyz",
+                typeCode: "z",
+                compress: true,
+                payload: pyzArchive({
+                  magic: PYTHON_39_MAGIC,
+                  modules: [{ name: "app", typeCode: 0, body: payload }],
+                }),
+              },
+            ],
+          });
+    await withPyInstallerReader(bytes, async (reader, handle) => {
+      const entered = deferred(),
+        release = deferred();
+      const caller = new AbortController();
+      const reason = new Error("cancel first archive parse");
+      const nativeRead = handle.read.bind(handle);
+      const spool = PyInstallerPyzStore.prototype.spool;
+      const images: FileHandle[] = [];
+      let held = false;
+      const pause = async () => {
+        if (held) return;
+        held = true;
+        entered.resolve();
+        await release.promise;
+      };
+      const reads = vi
+        .spyOn(handle, "read")
+        .mockImplementation(async (...args: unknown[]) => {
+          const result = await Reflect.apply(nativeRead, handle, args);
+          if (stage === "carrier") await pause();
+          return result;
+        });
+      const spools = vi
+        .spyOn(PyInstallerPyzStore.prototype, "spool")
+        .mockImplementation(async function (
+          this: PyInstallerPyzStore,
+          ...args
+        ) {
+          const result = await spool.apply(this, args);
+          images.push(result.handle);
+          if (stage === "private image") await pause();
+          return result;
+        });
+      const first = members(reader, caller.signal).then(
+        () => undefined,
+        (cause: unknown) => cause,
+      );
+      let second: Promise<ArtifactEntry[]> | undefined;
+      try {
+        await entered.promise;
+        expect(reader.formatMetadata()).toBeUndefined();
+        if (peer) {
+          second = members(reader);
+          void second.catch(() => undefined);
+        }
+        caller.abort(reason);
+        release.resolve();
+        expect(await first).toMatchObject({
+          reason: "cancelled",
+          cause: reason,
+        });
+        const retried = await (second ?? members(reader));
+        const selected = retried.find(
+          ({ path }) =>
+            path === (stage === "carrier" ? "data" : "PYZ.pyz/app.pyc"),
+        );
+        expect(selected).toBeDefined();
+        const content = await buffer(await reader.open(selected!));
+        expect(
+          content.subarray(stage === "carrier" ? 0 : 16).equals(payload),
+        ).toBe(true);
+        if (stage === "private image") {
+          expect(images).toHaveLength(2);
+          expect(images[0]!.fd).toBe(-1);
+        }
+      } finally {
+        release.resolve();
+        await first;
+        await second?.catch(() => undefined);
+        reads.mockRestore();
+        spools.mockRestore();
+      }
+    });
+  },
+);
+
+it("does not charge a cancelled namespace projection against the next parse", async () => {
+  const name = "a/".repeat(2000) + "x";
+  const names = [
+    name,
+    ...Array.from({ length: 256 }, (_, index) => `sibling-${index}`),
+  ];
+  const bytes = cArchiveExecutable({
+    layout: "2.1",
+    pythonVersion: 309,
+    members: names.map((name) => ({
+      name,
+      typeCode: "x",
+      payload: Buffer.alloc(0),
+      compress: false,
+    })),
+  });
+  await withPyInstallerReader(bytes, async (reader) => {
+    const caller = new AbortController();
+    const reason = new Error("cancel after namespace reservation");
+    let scheduled = false;
+    const normalize = artifactPaths.normalizeArtifactPath;
+    const paths = vi
+      .spyOn(artifactPaths, "normalizeArtifactPath")
+      .mockImplementation((path) => {
+        const result = normalize(path);
+        // The projection yields at its next batch after reserving this deep path.
+        if (path === name && !scheduled) {
+          scheduled = true;
+          globalThis.setImmediate(() => caller.abort(reason));
+        }
+        return result;
+      });
+    try {
+      await expect(members(reader, caller.signal)).rejects.toMatchObject({
+        reason: "cancelled",
+        cause: reason,
+      });
+      expect(reader.formatMetadata()).toBeUndefined();
+      const retried = await members(reader);
+      expect(retried.map(({ path }) => path)).toEqual(names);
+      expect(await buffer(await reader.open(retried[0]!))).toHaveLength(0);
+    } finally {
+      paths.mockRestore();
+    }
+  });
+});
+
+it("keeps a physical parse failure primary despite simultaneous caller cancellation", async () => {
+  await withPyInstallerReader(
+    archive(Buffer.from("body")),
+    async (reader, handle) => {
+      const entered = deferred(),
+        release = deferred();
+      const caller = new AbortController();
+      const io = new Error("carrier read EIO");
+      const reads = vi.spyOn(handle, "read").mockImplementation(async () => {
+        entered.resolve();
+        await release.promise;
+        caller.abort(new Error("caller also leaves"));
+        throw io;
+      });
+      const first = members(reader, caller.signal).catch(
+        (cause: unknown) => cause,
+      );
+      let second: Promise<unknown> | undefined;
+      try {
+        await entered.promise;
+        second = members(reader).catch((cause: unknown) => cause);
+        release.resolve();
+        const original = await first;
+        expect(original).toMatchObject({ reason: "io", cause: io });
+        expect(await second).toBe(original);
+        reads.mockRestore();
+        await expect(members(reader)).rejects.toBe(original);
+        expect(reader.formatMetadata()).toBeUndefined();
+      } finally {
+        release.resolve();
+        await first;
+        await second;
+        reads.mockRestore();
+      }
+    },
+  );
+});
+
+it("retains failed cancelled-attempt cleanup and retries its exact private owner on close", async () => {
+  const bytes = cArchiveExecutable({
+    layout: "2.1",
+    pythonVersion: 309,
+    members: [
+      {
+        name: "PYZ.pyz",
+        typeCode: "z",
+        compress: true,
+        payload: pyzArchive({ magic: PYTHON_39_MAGIC, modules: [] }),
+      },
+    ],
+  });
+  await withPyInstallerReader(bytes, async (reader) => {
+    const caller = new AbortController();
+    const reason = new Error("cancel private parse");
+    const cleanupFailure = new Error("private cancelled close denied");
+    const spool = PyInstallerPyzStore.prototype.spool;
+    const nativeClose = OwnedFileHandle.prototype.close;
+    let retained: FileHandle | undefined;
+    let attempts = 0;
+    const spools = vi
+      .spyOn(PyInstallerPyzStore.prototype, "spool")
+      .mockImplementation(async function (this: PyInstallerPyzStore, ...args) {
+        const result = await spool.apply(this, args);
+        retained = result.handle;
+        caller.abort(reason);
+        return result;
+      });
+    const closes = vi
+      .spyOn(OwnedFileHandle.prototype, "close")
+      .mockImplementation(async function (this: OwnedFileHandle) {
+        if (this.handle === retained && ++attempts === 1) throw cleanupFailure;
+        await nativeClose.call(this);
+      });
+    try {
+      const failure = await members(reader, caller.signal).catch(
+        (cause: unknown) => cause,
+      );
+      expect(failure).toMatchObject({
+        reason: "cancelled",
+        cancellation: { reason },
+        cleanup: { reason: expect.stringContaining(cleanupFailure.message) },
+      });
+      expect(retained?.fd).toBeGreaterThanOrEqual(0);
+      await expect(members(reader)).rejects.toBe(failure);
+      expect(spools).toHaveBeenCalledTimes(1);
+      await reader.close();
+      expect(attempts).toBe(2);
+      expect(retained?.fd).toBe(-1);
+    } finally {
+      spools.mockRestore();
+      closes.mockRestore();
+    }
+  });
+});