diff --git a/src/io/PipeReader.rs b/src/io/PipeReader.rs index 83f8c081e15b..ea4543f68634 100644 --- a/src/io/PipeReader.rs +++ b/src/io/PipeReader.rs @@ -162,7 +162,15 @@ struct ReadScratchClaim; impl ReadScratchClaim { fn try_claim() -> Option { - 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 diff --git a/test/js/workerd/html-rewriter.test.js b/test/js/workerd/html-rewriter.test.js index 73e494b0b28f..aaf7276e0649 100644 --- a/test/js/workerd/html-rewriter.test.js +++ b/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, `

${text}

`).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 = `

${otherText}

`; + const otherRewritten = Buffer.alloc((otherPiece.length + 6) * count, `

${otherText}

`).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, `

${text}

`).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.