http: clear the request stream drain callback under the buffer mutex

src/http/ThreadSafeStreamBuffer.rs

@@ -9,9 +9,13 @@ pub struct ThreadSafeStreamBuffer {
     pub(crate) mutex: Mutex,
     /// Intrusive atomic refcount. Starts at 2: 1 for main thread and 1 for http thread.
     pub(crate) ref_count: bun_ptr::ThreadSafeRefCount<ThreadSafeStreamBuffer>,
-    /// callback will be called passing the context for the http callback
-    /// this is used to report when the buffer is drained and only if end chunk was not sent/reported
-    pub(crate) callback: Option<Callback>,
+    /// Invoked by the HTTP thread (`report_drain`) once everything buffered
+    /// has been written out and the end chunk has not been sent yet.
+    ///
+    /// Guarded by `mutex`, like `buffer`: `report_drain` reads it under the
+    /// lock on the HTTP thread while the JS thread may be clearing it in
+    /// `clear_drain_callback` for a request that is still in flight.
+    callback: Option<Callback>,
 }
 
 pub struct Callback {
@@ -102,13 +106,20 @@ impl ThreadSafeStreamBuffer {
         self.callback = Some(Callback::init(callback, context));
     }
 
+    /// Main thread. The HTTP thread may still be flushing this buffer (and
+    /// about to `report_drain`) when the JS side tears the request down, so
+    /// the callback can only be cleared under the same lock that
+    /// `report_drain` reads it under; once this returns it can no longer fire.
     pub fn clear_drain_callback(&mut self) {
+        let _guard = self.mutex.lock_guard();
         self.callback = None;
     }
 
     /// This is exclusively called from the http thread.
-    /// Buffer should be acquired before calling this.
+    /// The caller must hold the lock (`acquire`/`lock`); it is what keeps
+    /// `callback` from being cleared out from under this read.
     pub(crate) fn report_drain(&self) {
+        debug_assert!(self.mutex.is_held_by_current_thread());
         if self.buffer.is_empty() {
             if let Some(callback) = &self.callback {
                 callback.call();

src/runtime/webcore/fetch/FetchTasklet.rs

@@ -386,9 +386,9 @@ impl FetchTasklet {
     /// shared with the HTTP thread (mutex-guarded internally).
     #[inline]
     pub(crate) fn stream_buffer_mut<'r>(&self) -> Option<&'r mut ThreadSafeStreamBuffer> {
-        // SAFETY: see doc comment — counted ref keeps pointee live; mutex
-        // inside `ThreadSafeStreamBuffer` serialises cross-thread `buffer`
-        // access, and `callback` is main-thread-only.
+        // SAFETY: see doc comment — counted ref keeps pointee live; the mutex
+        // inside `ThreadSafeStreamBuffer` serialises every cross-thread access
+        // (`buffer` and the drain callback alike).
         self.request_body_streaming_buffer
             .map(|p| unsafe { &mut *p.as_ptr() })
     }
@@ -450,8 +450,10 @@ impl FetchTasklet {
         }
         if let Some(buffer) = self.request_body_streaming_buffer.take() {
             // SAFETY: intrusive-refcounted heap allocation from `ThreadSafeStreamBuffer::new`;
-            // this side holds one of the two initial refs. Mutex guards cross-thread access
-            // to `buffer`, and `callback` is only touched on the main thread (here).
+            // this side holds one of the two initial refs. The HTTP thread may still own its
+            // ref and be flushing (`start_request_stream` gets here with the request in
+            // flight); `clear_drain_callback` takes the buffer's mutex, so after it returns
+            // the HTTP thread can no longer call back into this tasklet through the buffer.
             unsafe { (*buffer.as_ptr()).clear_drain_callback() };
             ThreadSafeStreamBuffer::deref(buffer);
         }

test/js/web/fetch/fetch-abort-stream-body.test.ts

@@ -79,6 +79,51 @@ test
     expect(exitCode).toBe(0);
   });
 
+// A direct stream's pull() runs synchronously inside start_request_stream, and
+// that only happens once the HTTP thread has sent the headers and asked for the
+// body. Writing and then throwing from it tears the request down (clear_sink)
+// while the HTTP thread is still flushing the bytes just written and reporting
+// the buffer drained, so the JS side clears the buffer's drain callback at the
+// same moment the HTTP thread reads it; both have to go through the buffer's
+// mutex. Every iteration has to reject with pull's own error, and clearing the
+// callback must not deadlock against the HTTP thread holding the buffer.
+test.concurrent(
+  "request body pull() that writes and then throws rejects the fetch while the upload is in flight",
+  async () => {
+    await using server = Bun.serve({
+      port: 0,
+      async fetch(req) {
+        // Only answer once the client has torn the upload down, so the rejection
+        // below can only come from pull()'s error, never from a response.
+        await req.arrayBuffer().catch(() => {});
+        return new Response("unreachable");
+      },
+    });
+
+    const iterations = 50;
+    // Several chunks over the sink's 16 KiB high water mark, so the HTTP thread
+    // is woken and has something to flush (and report drained) while pull()
+    // throws on the JS thread.
+    const chunk = Buffer.alloc(64 * 1024, "x");
+    let pulls = 0;
+
+    for (let i = 0; i < iterations; i++) {
+      const error = new Error(`pull ${i}`);
+      const body = new ReadableStream({
+        type: "direct",
+        pull(controller) {
+          pulls++;
+          for (let j = 0; j < 4; j++) controller.write(chunk);
+          throw error;
+        },
+      });
+      await expect(fetch(server.url, { method: "POST", body })).rejects.toBe(error);
+    }
+
+    expect(pulls).toBe(iterations);
+  },
+);
+
 test("aborting fetch with a ReadableStream request body does not double-cancel the sink", async () => {
   await using proc = Bun.spawn({
     cmd: [bunExe(), join(import.meta.dir, "fetch-abort-stream-body-fixture.ts")],