ByteStream: re-signal a backpressured RewriterPipe after native-sink wiring writes the buffered bytes

src/runtime/api/html_rewriter.rs

@@ -886,8 +886,14 @@ impl RewriterPipe {
     #[inline]
     fn output_backpressured(&self) -> bool {
         if let Some(out) = self.output.get() {
-            return out.sink_paused.get()
-                || out.buffer.get().len() as BlobSizeType > self.high_water_mark.get();
+            if out.sink_paused.get() {
+                return true;
+            }
+            // A body-mixin collector grows `buffer` deliberately until Done.
+            if out.buffer_action.get().is_some() {
+                return false;
+            }
+            return out.buffer.get().len() as BlobSizeType > self.high_water_mark.get();
         }
         // No output ByteStream yet: the pre-stream buffer has no drain signal
         // (only `on_start_streaming`/`finish` consume it), so backpressuring

src/runtime/server/RequestContext.rs

@@ -3279,7 +3279,7 @@ 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 mut response_buf = byte_stream.take_buffer();
                             let buffer = response_buf.move_to_list();
                             let has_body_bytes = !buffer.is_empty();
                             this.response_buf_owned.set(buffer);
@@ -3308,6 +3308,8 @@ where
                                     this.as_ctx_ptr(),
                                 );
                             }
+                            // Wake the producer after the older bytes are queued.
+                            byte_stream.signal_drained();
                             return;
                         }
                     }

src/runtime/webcore/ByteStream.rs

@@ -223,10 +223,17 @@ impl ByteStream {
     }
 
     #[inline]
-    fn signal_drained(&self) {
+    pub(crate) fn signal_drained(&self) {
         self.parent_const().producer.get().ready(None, None);
     }
 
+    /// Take the buffered bytes without signalling the producer; the caller
+    /// writes them to the sink before [`Self::signal_drained`].
+    pub(crate) fn take_buffer(&self) -> Vec<u8> {
+        self.offset.set(0);
+        Vec::<u8>::move_from_list(self.buffer.replace(Vec::new()))
+    }
+
     /// Called by native fast-paths after wiring `self.sink`. Restores
     /// producer-side backpressure if it was already dropped (BufferAll).
     pub fn signal_consumer_attached(&self) {
@@ -743,10 +750,12 @@ impl ByteStream {
             return Ok(blob.to_promise(global_this, action)?);
         }
 
-        self.signal_drained();
         self.buffer_action
             .set(Some(BufferAction::new(action, global_this)));
-
-        Ok(self.buffer_action.get().as_ref().unwrap().value())
+        let promise = self.buffer_action.get().as_ref().unwrap().value();
+        // Signal after the action is installed so a backpressure-gated
+        // producer observes it; a synchronous producer may fulfil it inline.
+        self.signal_drained();
+        Ok(promise)
     }
 }

src/runtime/webcore/FileSink.rs

@@ -1557,11 +1557,15 @@ impl FileSink {
             |src| self.source.set(src),
         ) {
             readable_stream::NativeWireResult::Wired => {
-                self.writer
-                    .with_mut(|w| w.enable_keeping_process_alive(self.io_evtloop()));
-                if !self.must_be_kept_alive_until_eof.get() {
-                    self.must_be_kept_alive_until_eof.set(true);
-                    self.ref_();
+                // A synchronous producer may have driven `end_from_stream`
+                // (clears `source`) inline; no keepalive then.
+                if !matches!(self.source.get(), streams::SourceHandle::None) {
+                    self.writer
+                        .with_mut(|w| w.enable_keeping_process_alive(self.io_evtloop()));
+                    if !self.must_be_kept_alive_until_eof.get() {
+                        self.must_be_kept_alive_until_eof.set(true);
+                        self.ref_();
+                    }
                 }
                 return JSValue::UNDEFINED;
             }

src/runtime/webcore/ReadableStream.rs

@@ -337,10 +337,10 @@ impl ReadableStream {
                     return NativeWireResult::EndedInline(Some(err));
                 }
 
-                let buffered = byte_stream.drain();
-                let has_last = byte_stream.has_received_last_chunk.get();
+                let buffered = byte_stream.take_buffer();
+                let had_last = byte_stream.has_received_last_chunk.get();
                 if !buffered.is_empty() {
-                    let chunk = if has_last {
+                    let chunk = if had_last {
                         StreamResult::OwnedAndDone(buffered)
                     } else {
                         StreamResult::Owned(buffered)
@@ -354,10 +354,15 @@ impl ReadableStream {
                         _ => {}
                     }
                 }
-                if has_last {
+                if had_last {
                     byte_stream.sink.set(SinkHandle::None);
                     return NativeWireResult::EndedInline(None);
                 }
+                // Wake the producer after the older bytes are in the sink;
+                // any synchronous output routes through `on_data`.
+                if !byte_stream.sink_paused.get() {
+                    byte_stream.signal_drained();
+                }
                 return NativeWireResult::Wired;
             }
         }

src/runtime/webcore/s3/client.rs

@@ -1037,10 +1037,10 @@ 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();
+            let buffered = byte_stream.take_buffer();
+            let had_last = byte_stream.has_received_last_chunk.get();
             if !buffered.is_empty() {
-                let chunk = if has_last {
+                let chunk = if had_last {
                     crate::webcore::streams::StreamResult::OwnedAndDone(buffered)
                 } else {
                     crate::webcore::streams::StreamResult::Owned(buffered)
@@ -1062,15 +1062,18 @@ pub(crate) fn upload_stream(
                     _ => {}
                 }
             }
-            if has_last {
+            if had_last {
                 byte_stream.sink.set(crate::webcore::SinkHandle::None);
                 sink.source.clear();
                 if !sink.ended {
                     let _ = sink.end(None);
                 }
                 ctx.handle_resolve_stream();
+            } else if !byte_stream.sink_paused.get() {
+                // Wake the producer after the older bytes are in the sink.
+                byte_stream.signal_drained();
             }
-            // `!has_last`: the stream-pump +1 (rc=2) is released by
+            // `!had_last`: the stream-pump +1 (rc=2) is released by
             // `NetworkSink::end_from_stream` after the terminal write/fail so the
             // sink outlives the synchronous `resolve()` re-entry.
             return Ok(end_promise_value);

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

@@ -1708,6 +1708,160 @@ const request_types = ["/", "/gzip", "/chunked/gzip", "/chunked", "/file", "/fil
   });
 });
 
+// An async handler suspends mid-input; once resumed it emits more output than
+// RewriterPipe's high-water mark (16 KiB) into the already-realised ByteStream
+// and parks on output backpressure with `input_ended` set. Wiring a native sink
+// afterwards must re-signal the rewriter once the buffered bytes are in the
+// sink so `end_rewrite` runs.
+describe("output ByteStream backpressured when a native sink is wired", () => {
+  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;
+  const out = prefix + "<x>!</x>" + suffix;
+
+  // Return a ReadableStream whose ByteStream buffer already holds the full
+  // rewritten output while the rewriter itself is still parked on
+  // `output_backpressured()` with `input_ended = true`.
+  async function makeParked() {
+    const gate = Promise.withResolvers();
+    const res = new HTMLRewriter()
+      .on("x", {
+        element: e => {
+          e.setInnerContent("!");
+          return gate.promise;
+        },
+      })
+      .transform(new Response(html));
+    const body = res.body;
+    gate.resolve();
+    // Let the handler-promise reaction run so `resume_rewrite` emits `suffix`
+    // into the ByteStream buffer.
+    await 0;
+    return { body, res };
+  }
+
+  it("completes when wired to a second HTMLRewriter", async () => {
+    const { body } = await makeParked();
+    const text = await new HTMLRewriter().transform(new Response(body)).text();
+    expect(text).toBe(out);
+  });
+
+  it("completes when wired to a second HTMLRewriter via the Response", async () => {
+    const { res } = await makeParked();
+    const text = await new HTMLRewriter().transform(res).text();
+    expect(text).toBe(out);
+  });
+
+  it("delivers bytes in order", async () => {
+    const { body } = await makeParked();
+    let seen = "";
+    await new HTMLRewriter()
+      .onDocument({ text: t => void (seen += t.text) })
+      .transform(new Response(body))
+      .text();
+    expect(seen).toBe(prefix.slice(3, -4) + "!" + suffix.slice(3, -4));
+  });
+
+  it("completes when returned from Bun.serve", async () => {
+    await using server = Bun.serve({
+      port: 0,
+      fetch: async () => {
+        const { body } = await makeParked();
+        return new Response(body);
+      },
+    });
+    const text = await (await fetch(server.url)).text();
+    expect(text).toBe(out);
+  });
+
+  it("completes when the Response is returned from Bun.serve", async () => {
+    await using server = Bun.serve({
+      port: 0,
+      fetch: async () => (await makeParked()).res,
+    });
+    const text = await (await fetch(server.url)).text();
+    expect(text).toBe(out);
+  });
+
+  it("completes when read via .arrayBuffer()", async () => {
+    const { body } = await makeParked();
+    const buf = await new Response(body).arrayBuffer();
+    expect(Buffer.from(buf).toString()).toBe(out);
+  });
+
+  it("completes when read via .text()", async () => {
+    const { body } = await makeParked();
+    expect(await new Response(body).text()).toBe(out);
+  });
+
+  it("completes when used as Bun.spawn stdin", async () => {
+    const { body } = await makeParked();
+    await using proc = Bun.spawn({
+      cmd: [bunExe(), "-e", "process.stdout.write(await Bun.stdin.text())"],
+      env: bunEnv,
+      stdin: body,
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    expect(stderr).toBe("");
+    expect(stdout).toBe(out);
+    expect(exitCode).toBe(0);
+  });
+
+  // Streaming input, `input_ended=false`: the handler itself emits the >16 KiB
+  // output chunk, and the direct-stream's `pull()` is parked on `flush(true)`
+  // with more bytes to write. `signal_drained()` must wake the rewriter into
+  // `drain_pending_input()`'s upstream-`src.ready()` + `pending.run()` tail
+  // so the direct stream's second `c.write()` is delivered.
+  const streamed = prefix + "<x>" + suffix.slice(3, -4) + "</x>" + suffix;
+  async function makeParkedStreaming() {
+    const gate = Promise.withResolvers();
+    const input = new ReadableStream({
+      type: "direct",
+      async pull(c) {
+        c.write("<x></x>");
+        await c.flush(true);
+        c.write(suffix);
+        c.close();
+      },
+    });
+    const res = new HTMLRewriter()
+      .on("x", {
+        element: e => {
+          e.before(prefix, { html: true });
+          e.setInnerContent(suffix.slice(3, -4));
+          return gate.promise;
+        },
+      })
+      .transform(new Response(input));
+    const body = res.body;
+    gate.resolve();
+    await 0;
+    return { body };
+  }
+
+  it("completes with a streaming input (second HTMLRewriter)", async () => {
+    const { body } = await makeParkedStreaming();
+    const text = await new HTMLRewriter().transform(new Response(body)).text();
+    expect(text).toBe(streamed);
+  });
+
+  it("completes with a streaming input (Bun.serve)", async () => {
+    await using server = Bun.serve({
+      port: 0,
+      fetch: async () => new Response((await makeParkedStreaming()).body),
+    });
+    const text = await (await fetch(server.url)).text();
+    expect(text).toBe(streamed);
+  });
+
+  it("completes with a streaming input (.text())", async () => {
+    const { body } = await makeParkedStreaming();
+    expect(await new Response(body).text()).toBe(streamed);
+  });
+});
+
 const payloads = [
   {
     name: "direct",