PipeReader: keep the outer read loop's scratch claim when a nested read is refused

src/io/PipeReader.rs

@@ -162,7 +162,15 @@ struct ReadScratchClaim;
 
 impl ReadScratchClaim {
     fn try_claim() -> Option<Self> {
-        READ_SCRATCH_IN_USE.with(|in_use| (!in_use.replace(true)).then_some(Self))
+        READ_SCRATCH_IN_USE.with(|in_use| {
+            if in_use.get() {
+                // Not `then_some(Self)`: its argument is built before the test,
+                // and dropping that refused claim would release the outer one.
+                return None;
+            }
+            in_use.set(true);
+            Some(Self)
+        })
     }
 }
 
@@ -902,8 +910,14 @@ impl PosixBufferedReader {
             } else {
                 // SAFETY: caller contract; `maxbuf` is Copy, borrow ends at `;`.
                 let maxbuf = unsafe { (*this).maxbuf };
+                // Streaming delivers and clears per read, so this reserve is the
+                // chunk size. Nested reads (a consumer re-pulling from inside its
+                // chunk handler, e.g. `process.stdin.on("data")`) all land here
+                // rather than in the scratch, so give them a whole default pipe
+                // buffer per read.
+                let reserve = if streaming { 64 * 1024 } else { 16 * 1024 };
                 // SAFETY: caller contract; borrow ends at `;`.
-                unsafe { (*this)._buffer.reserve(16 * 1024) };
+                unsafe { (*this)._buffer.reserve(reserve) };
                 // SAFETY: caller contract. `sys::read_nonblocking` writes only
                 // initialized bytes into the prefix it reports; `commit_spare`
                 // exposes exactly that prefix. The `_buffer` borrow ends before

test/js/workerd/html-rewriter.test.js

@@ -1895,8 +1895,18 @@ describe("streamed input pacing", () => {
   const count = 2500; // ~2.4 MB of input: many upstream chunks
   const input = Buffer.alloc(piece.length * count, piece).toString();
   const rewritten = Buffer.alloc((piece.length + 6) * count, `<p x="1">${text}</p>`).toString();
-  const dir = tempDirWithFiles("hr-pacing", { "in.html": input });
+  // A second document made of different bytes, for the tests below that read
+  // it while another document is being parsed: if that read lands in the
+  // other document's buffer, these bytes show up in its output.
+  const otherText = Buffer.alloc(1000, "b").toString();
+  const otherPiece = `<p>${otherText}</p>`;
+  const otherRewritten = Buffer.alloc((otherPiece.length + 6) * count, `<p x="1">${otherText}</p>`).toString();
+  const dir = tempDirWithFiles("hr-pacing", {
+    "in.html": input,
+    "other.html": Buffer.alloc(otherPiece.length * count, otherPiece).toString(),
+  });
   const file = path.join(dir, "in.html");
+  const otherFile = path.join(dir, "other.html");
 
   function transformInput(body = Bun.file(file)) {
     let seen = 0;
@@ -2077,6 +2087,62 @@ describe("streamed input pacing", () => {
     expect(await inner).toBe(rewritten);
   });
 
+  // Starting a transform inside the handler and reading it are two reads
+  // nested in the outer one, which still has its document in the read buffer:
+  // the second nested read must be kept out of it just like the first.
+  it("a handler may transform and read another file", async () => {
+    let inner;
+    const outer = new HTMLRewriter()
+      .on("p", {
+        element(e) {
+          e.setAttribute("x", "1");
+          inner ??= transformInput(Bun.file(otherFile)).res.text();
+        },
+      })
+      .transform(new Response(Bun.file(file)));
+    expect(await outer.text()).toBe(rewritten);
+    expect(await inner).toBe(otherRewritten);
+  });
+
+  // Same, with the outer document arriving on a pipe (the child's stdin): a
+  // pipe is read by a different loop than a regular file, out of the same
+  // buffer.
+  it("a handler may transform and read another file while its document arrives on stdin", async () => {
+    const pieces = 64;
+    await using proc = Bun.spawn({
+      cmd: [
+        bunExe(),
+        "-e",
+        `const attr = { element: e => void e.setAttribute("x", "1") };
+         let inner;
+         const outer = new HTMLRewriter()
+           .on("p", {
+             element(e) {
+               attr.element(e);
+               inner ??= new HTMLRewriter().on("p", attr).transform(new Response(Bun.file(process.argv[1]))).text();
+             },
+           })
+           .transform(new Response(Bun.stdin));
+         const out = await outer.text();
+         console.log(JSON.stringify({ out, innerLength: (await inner).length }));`,
+        otherFile,
+      ],
+      env: bunEnv,
+      stdin: "pipe",
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    proc.stdin.write(Buffer.alloc(piece.length * pieces, piece));
+    proc.stdin.end();
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    expect(stderr).toBe("");
+    expect(JSON.parse(stdout)).toEqual({
+      out: Buffer.alloc((piece.length + 6) * pieces, `<p x="1">${text}</p>`).toString(),
+      innerLength: otherRewritten.length,
+    });
+    expect(exitCode).toBe(0);
+  });
+
   // Regular-file reads are synchronous on POSIX, so reading ahead of the
   // consumer would show up as `transform()` itself reading (and buffering the
   // rewrite of) the entire file. Sparse files keep these cheap.