ByteStream::drain: empty buffer before signal_drained (JS-reader path)

src/runtime/server/RequestContext.rs

@@ -3279,17 +3279,16 @@ where
                                 .set(readable_stream::Strong::init(stream, global_this));
 
                             this.byte_stream.set(Some(byte_stream_nn));
-                            let mut response_buf = byte_stream.drain();
-                            let buffer = response_buf.move_to_list();
-                            let has_body_bytes = !buffer.is_empty();
-                            this.response_buf_owned.set(buffer);
 
                             // we don't set size here because even if we have a hint
                             // uWebSockets won't let us partially write streaming content
                             this.blob.with_mut(|b| b.detach());
 
                             // if we've received metadata and part of the body, send everything we can and drain
-                            if has_body_bytes {
+                            if !byte_stream.buffer.get().is_empty() {
+                                let mut buf = byte_stream.buffer.replace(Vec::new());
+                                byte_stream.offset.set(0);
+                                this.response_buf_owned.set(buf.move_to_list());
                                 resp.run_corked_with_type(
                                     Self::drain_response_buffer_and_metadata_corked,
                                     this.as_ctx_ptr(),
@@ -3308,6 +3307,10 @@ where
                                     this.as_ctx_ptr(),
                                 );
                             }
+                            // Wake the producer now that the buffered bytes are
+                            // written; any output it emits synchronously reaches
+                            // `write_chunk` after them.
+                            byte_stream.flush_to_sink();
                             return;
                         }
                     }

src/runtime/webcore/ByteStream.rs

@@ -167,7 +167,17 @@ impl ByteStream {
             return;
         }
         self.sink_paused.set(false);
+        self.flush_to_sink();
+    }
 
+    /// Write any buffered bytes to `self.sink`, then [`signal_drained`]. The
+    /// native-sink wiring paths call this (instead of [`drain`] + a manual
+    /// `sink.write`) so the buffered bytes reach the sink before the producer
+    /// is woken; a producer whose `on_ready` feeds synchronously (RewriterPipe)
+    /// would otherwise emit newer bytes to the already-installed sink ahead of
+    /// the ones the caller is still holding. Clears `self.sink` when the write
+    /// or the last-chunk flag ends the stream.
+    pub(crate) fn flush_to_sink(&self) {
         let sink = *self.sink.get();
         if sink.is_none() {
             return;
@@ -672,12 +682,24 @@ impl ByteStream {
         // deallocate the storage backing `&mut self` (dangling UAF).
     }
 
+    /// JS-reader drain (native-source adapter's `handle.drain()` and the
+    /// `has_received_last_chunk` blob hand-off). Callers that have already
+    /// installed `self.sink` must use [`flush_to_sink`] instead so the bytes
+    /// are written before the producer is woken.
     pub(crate) fn drain(&self) -> Vec<u8> {
-        if !self.buffer.get().is_empty() {
+        if self.buffer.get().is_empty() {
+            return Vec::<u8>::default();
+        }
+        let drained = Vec::<u8>::move_from_list(self.buffer.replace(Vec::new()));
+        // `materializeNativeSource` can reach here with a sink already wired
+        // (getReader() on a natively-locked stream materialises before the
+        // lock check throws); signalling then would let a synchronous producer
+        // emit newer bytes to that sink ahead of `drained`, so only signal on
+        // the intended JS-reader path.
+        if self.sink.get().is_none() {
             self.signal_drained();
-            return Vec::<u8>::move_from_list(self.buffer.replace(Vec::new()));
         }
-        Vec::<u8>::default()
+        drained
     }
 
     /// Take a pre-attach `StreamResult::Err` stashed by [`Self::append`].

src/runtime/webcore/ReadableStream.rs

@@ -320,7 +320,7 @@ impl ReadableStream {
         owner_cell: JSValue,
         set_source: impl FnOnce(streams::SourceHandle),
     ) -> NativeWireResult {
-        use streams::{SourceHandle, Start, StreamError, StreamResult, Writable};
+        use streams::{SourceHandle, Start, StreamError, StreamResult};
         use webcore::SinkHandle;
 
         if let Some(byte_stream) = self.ptr.bytes() {
@@ -337,25 +337,8 @@ impl ReadableStream {
                     return NativeWireResult::EndedInline(Some(err));
                 }
 
-                let buffered = byte_stream.drain();
-                let has_last = byte_stream.has_received_last_chunk.get();
-                if !buffered.is_empty() {
-                    let chunk = if has_last {
-                        StreamResult::OwnedAndDone(buffered)
-                    } else {
-                        StreamResult::Owned(buffered)
-                    };
-                    match sink.write(&chunk) {
-                        Writable::Backpressure(_) => byte_stream.sink_paused.set(true),
-                        Writable::Done | Writable::Err(_) => {
-                            byte_stream.sink.set(SinkHandle::None);
-                            return NativeWireResult::EndedInline(None);
-                        }
-                        _ => {}
-                    }
-                }
-                if has_last {
-                    byte_stream.sink.set(SinkHandle::None);
+                byte_stream.flush_to_sink();
+                if byte_stream.sink.get().is_none() {
                     return NativeWireResult::EndedInline(None);
                 }
                 return NativeWireResult::Wired;

src/runtime/webcore/s3/client.rs

@@ -1037,33 +1037,8 @@ pub(crate) fn upload_stream(
                 return Ok(end_promise_value);
             }
 
-            let buffered = byte_stream.drain();
-            let has_last = byte_stream.has_received_last_chunk.get();
-            if !buffered.is_empty() {
-                let chunk = if has_last {
-                    crate::webcore::streams::StreamResult::OwnedAndDone(buffered)
-                } else {
-                    crate::webcore::streams::StreamResult::Owned(buffered)
-                };
-                match sink.write(&chunk) {
-                    crate::webcore::streams::Writable::Backpressure(_) => {
-                        byte_stream.sink_paused.set(true);
-                    }
-                    crate::webcore::streams::Writable::Done
-                    | crate::webcore::streams::Writable::Err(_) => {
-                        byte_stream.sink.set(crate::webcore::SinkHandle::None);
-                        sink.source.clear();
-                        if !sink.ended {
-                            let _ = sink.end(None);
-                        }
-                        ctx.handle_resolve_stream();
-                        return Ok(end_promise_value);
-                    }
-                    _ => {}
-                }
-            }
-            if has_last {
-                byte_stream.sink.set(crate::webcore::SinkHandle::None);
+            byte_stream.flush_to_sink();
+            if byte_stream.sink.get().is_none() {
                 sink.source.clear();
                 if !sink.ended {
                     let _ = sink.end(None);

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

@@ -1741,6 +1741,126 @@ 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 = "<p>abc</p>";
+  const chunk = Buffer.alloc(50 * unit.length, unit).toString();
+  const expected = Buffer.alloc(50 * 16, '<p x="1">abc</p>').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);
+  });
+});
+
+// The same output-backpressure state, consumed by a native sink instead of a
+// JS reader: the sink is installed on the output ByteStream while it still
+// holds more than the rewriter's high-water mark in buffered bytes, and the
+// rewriter's input is already fully consumed (input_ended). flush_to_sink()
+// must write the buffered bytes to the sink before waking the rewriter (which
+// then runs end_rewrite() synchronously).
+describe("output consumed by a native sink after output backpressure", () => {
+  const prefix = "<p>" + Buffer.alloc(100, "A").toString() + "</p>";
+  const suffix = "<q>" + Buffer.alloc(30_000, "B").toString() + "</q>";
+  const html = prefix + "<x></x>" + suffix;
+
+  // Produce a rewriter output ByteStream in the "input_ended + buffer > hwm"
+  // state: .body materialises the ByteStream; resolving the handler gate lets
+  // resume_rewrite() push the 30 KB suffix into that buffer; drain_pending_input
+  // then returns early on output_backpressured(), leaving end_rewrite() owed.
+  async function makeParked() {
+    const gate = Promise.withResolvers();
+    const out = new HTMLRewriter().on("x", { element: () => gate.promise }).transform(new Response(html));
+    const body = out.body;
+    gate.resolve();
+    await 0;
+    return { out, body };
+  }
+
+  it("chained HTMLRewriter receives the full output in order", async () => {
+    const { body } = await makeParked();
+    const chained = new HTMLRewriter().transform(new Response(body));
+    expect(await chained.text()).toBe(html);
+  });
+
+  it("Bun.serve returning the rewriter Response sends the full output", async () => {
+    await using server = Bun.serve({
+      port: 0,
+      async fetch() {
+        return (await makeParked()).out;
+      },
+    });
+    expect(await (await fetch(server.url)).text()).toBe(html);
+  });
+
+  it("fetch request body wired to the output ByteStream uploads the full output", async () => {
+    const received = Promise.withResolvers();
+    await using server = Bun.serve({
+      port: 0,
+      async fetch(req) {
+        received.resolve(await req.text());
+        return new Response("ok");
+      },
+    });
+    const { body } = await makeParked();
+    await fetch(server.url, { method: "POST", body });
+    expect(await received.promise).toBe(html);
+  });
+});
+
 payloads.forEach(type => {
   type.test(`works with payload of type ${type.name}`, async () => {
     let calls = 0;