diff --git a/src/runtime/webcore/ByteStream.rs b/src/runtime/webcore/ByteStream.rs index f873f3b231b0..189d49f752ed 100644 --- a/src/runtime/webcore/ByteStream.rs +++ b/src/runtime/webcore/ByteStream.rs @@ -673,11 +673,24 @@ impl ByteStream { } pub(crate) fn drain(&self) -> Vec { - if !self.buffer.get().is_empty() { + if self.buffer.get().is_empty() { + return Vec::::default(); + } + // Empty first so a producer.on_ready that checks `buffer.len()` sees it + // drained. RewriterPipe is the one producer whose on_ready feeds + // synchronously; when a sink is installed the native-sink wiring caller + // writes `drained` to it afterward, so waking the rewriter here would + // let it reach the sink first. Async producers only schedule work. + let drained = Vec::::move_from_list(self.buffer.replace(Vec::new())); + let sync_behind_sink = self.sink.get().is_some() + && matches!( + self.parent_const().producer.get(), + streams::SourceHandle::HTMLRewriter(_) + ); + if !sync_behind_sink { self.signal_drained(); - return Vec::::move_from_list(self.buffer.replace(Vec::new())); } - Vec::::default() + drained } /// Take a pre-attach `StreamResult::Err` stashed by [`Self::append`]. diff --git a/test/js/workerd/html-rewriter.test.js b/test/js/workerd/html-rewriter.test.js index a6380cd1e85f..f5edd7771941 100644 --- a/test/js/workerd/html-rewriter.test.js +++ b/test/js/workerd/html-rewriter.test.js @@ -1741,6 +1741,71 @@ const payloads = [ }, ]; +// TextEncoderStream → HTMLRewriter → CompressionStream: the rewriter's input is +// driven by the native-sink bypass (encodeIntoSink), and its output ByteStream +// is consumed by a JS reader (pipeThrough), so its `drain()` is what wakes the +// rewriter once a batch of output bytes has been taken. With enough elements +// that one input chunk's output exceeds the rewriter's output high-water mark, +// the rewriter's `write` returns Backpressure; this must eventually resolve. +describe("output consumed via pipeThrough after a native-sink input transform", () => { + // ~50 elements → ~800 bytes of rewriter output per input chunk, well above + // the default output high-water mark (256). + const unit = "

abc

"; + const chunk = Buffer.alloc(50 * unit.length, unit).toString(); + const expected = Buffer.alloc(50 * 16, '

abc

').toString(); + + // The pull() macrotask yield is load-bearing: a synchronous pull lets the + // rewrite drain into the pre-stream output_buffer before the downstream + // reader attaches, so the ByteStream drain() path under test is never hit. + function makeRewritten(chunks = [chunk]) { + let i = 0; + const body = new ReadableStream({ + async pull(c) { + await Bun.sleep(0); + if (i < chunks.length) c.enqueue(chunks[i++]); + else c.close(); + }, + }); + return new HTMLRewriter() + .on("p", { element: e => e.setAttribute("x", "1") }) + .transform(new Response(body.pipeThrough(new TextEncoderStream()))); + } + + it("completes when read with arrayBuffer()", async () => { + const compressed = makeRewritten().body.pipeThrough(new CompressionStream("gzip")); + const buf = await new Response(compressed).arrayBuffer(); + const text = await new Response(new Blob([buf]).stream().pipeThrough(new DecompressionStream("gzip"))).text(); + expect(text).toBe(expected); + }); + + it("completes when served over HTTP", async () => { + await using server = Bun.serve({ + port: 0, + fetch: () => new Response(makeRewritten().body.pipeThrough(new CompressionStream("gzip"))), + }); + const buf = await (await fetch(server.url)).arrayBuffer(); + const text = await new Response(new Blob([buf]).stream().pipeThrough(new DecompressionStream("gzip"))).text(); + expect(text).toBe(expected); + }); + + // CompressionStream is the only native-byte-transform that can sit downstream + // of the rewriter's ByteStream output; cover each format so a later per-codec + // regression shows up here. + it.each(["gzip", "deflate", "deflate-raw", "brotli", "zstd"])("completes for CompressionStream(%s)", async format => { + const compressed = makeRewritten().body.pipeThrough(new CompressionStream(format)); + const text = await new Response(compressed.pipeThrough(new DecompressionStream(format))).text(); + expect(text).toBe(expected); + }); + + // Two input chunks → the second chunk is what the Backpressure wake must + // re-pull from upstream (the single-chunk case only owes the end() call). + it("completes across multiple input chunks", async () => { + const compressed = makeRewritten([chunk, chunk]).body.pipeThrough(new CompressionStream("gzip")); + const text = await new Response(compressed.pipeThrough(new DecompressionStream("gzip"))).text(); + expect(text).toBe(expected + expected); + }); +}); + payloads.forEach(type => { type.test(`works with payload of type ${type.name}`, async () => { let calls = 0;