Bun.serve: run the microtask checkpoint when a writable callback resumes script

src/runtime/server/RequestContext.rs

@@ -1730,6 +1730,10 @@ where
         let pinned = RequestContextRef::pin(this);
         let this = pinned.ctx();
         if let Some(sink) = this.sink_mut() {
+            // The drain wakes parked script. Its jobs run when the scope ends, inside the pin.
+            let _entered = sink
+                .on_writable_reaches_script()
+                .then(|| this.server().vm().enter_event_loop_scope());
             return sink.on_writable(write_offset, _resp);
         }
         true
@@ -3393,15 +3397,20 @@ where
         _resp: uws::AnyResponse,
     ) -> bool {
         ctx_log!("onWritableByteStream");
-        // SAFETY: `this` is live for the callback; only a shared view is
-        // formed, so the `resume()` re-entry into `write_chunk` is fine.
+        // SAFETY: `this` is live until `resume()`, which may release it (`end_chunk`): nothing after
+        // `resume()` reads it. Only a shared view is formed, so the re-entry into `write_chunk` is fine.
         let this = unsafe { &*this };
         debug_assert!(this.resp.get().is_some());
         if this.is_aborted_or_ended() {
             return false;
         }
         if let Some(bs) = this.byte_stream.get() {
-            bun_ptr::BackRef::from(bs).resume();
+            let bs = bun_ptr::BackRef::from(bs);
+            // A rewriter's handlers and its input's `pull()` run in `resume()`, their jobs when the scope ends.
+            let _entered = bs
+                .producer_runs_script()
+                .then(|| this.server().vm().enter_event_loop_scope());
+            bs.resume();
         }
         true
     }

src/runtime/webcore/ByteStream.rs

@@ -558,6 +558,25 @@ impl ByteStream {
         self.parent_const().producer.get().ready(None, None);
     }
 
+    /// Whether the producer answers [`Self::signal_drained`] with script, not only with bytes.
+    #[inline]
+    pub(crate) fn producer_runs_script(&self) -> bool {
+        match self.parent_const().producer.get() {
+            streams::SourceHandle::HTMLRewriter(_) => true,
+            streams::SourceHandle::None
+            | streams::SourceHandle::FetchResponseBody(_)
+            | streams::SourceHandle::ServerRequestBody(_)
+            | streams::SourceHandle::S3DownloadBody(_)
+            | streams::SourceHandle::TestingCancelOnDrain(_)
+            // The source of a sink, never the producer of a `ByteStream`.
+            | streams::SourceHandle::JSController(_)
+            | streams::SourceHandle::ByteStream(_)
+            | streams::SourceHandle::FileReader(_)
+            | streams::SourceHandle::Subprocess(_)
+            | streams::SourceHandle::ShellWritable(_) => false,
+        }
+    }
+
     /// Take the untaken bytes. The caller calls [`Self::signal_drained`] after it writes them.
     pub(crate) fn take_buffer(&self) -> Vec<u8> {
         Vec::<u8>::move_from_list(self.buffered.take())

src/runtime/webcore/streams.rs

@@ -1343,6 +1343,11 @@ impl<const SSL: bool> HTTPServerWritable<SSL> {
         &self.buffer[self.offset as usize..]
     }
 
+    /// Whether script waits for the drain: the pump or an awaited `write()`, `flush(true)`, `end()`.
+    pub(crate) fn on_writable_reaches_script(&self) -> bool {
+        self.source_pending_pull || self.pending_flush().is_some()
+    }
+
     pub(crate) fn on_writable(&mut self, write_offset: u64, _res: uws::AnyResponse) -> bool {
         // write_offset is the amount of data that was written not how much we need to write
         bun_core::scoped_log!(HTTPServerWritableLog, "onWritable ({})", write_offset);