Cancel the JS source when a ReadableStream pump is dropped mid-stream

docs/streams.md

@@ -88,7 +88,8 @@ and a **queue of pending reads**. The controller uses four callback functions
 
 - **start** -- invoked immediately when the `ReadableStream` is created
 - **pull** -- invoked to request more data from the source
-- **cancel** -- invoked when the stream is explicitly canceled
+- **cancel** -- invoked when the stream is canceled, either explicitly via `cancel()` or
+  when its consumer goes away (e.g. the client disconnects from a streaming response body)
 - **size** -- determines the size of a chunk for backpressure calculations
 
 When the stream is created, the start algorithm runs immediately. Once it completes,
@@ -413,7 +414,10 @@ these APIs were built for Internal streams and use kj async I/O internally.
 To bridge this, Standard `ReadableStream`s can be consumed via the `ReadableStreamSource`
 API (the same API Internal streams use). When `pumpTo()` is called on the adapter, it
 acquires the isolate lock and runs a promise loop: read from the JS stream, write to the
-kj output, repeat until the data is exhausted or an error occurs.
+kj output, repeat until the data is exhausted or an error occurs. If the client disconnects
+before then and the response body is dropped, the JS-controller pump
+(`ReadableStreamJsController::pumpTo`) schedules the stream's cancel algorithm so the
+source stops producing data.
 
 ## The Complexity Budget
 

src/workerd/api/streams-test.c++

@@ -198,5 +198,152 @@ KJ_TEST("ReadableStream pumpTo pending write cancellation regression") {
   KJ_ASSERT(events[2] == "sink was destroyed");
 }
 
+KJ_TEST("ReadableStream pumpTo cancels the JS source when dropped mid-stream") {
+  // Regression test for https://github.com/cloudflare/workerd/issues/6832.
+  //
+  // When the pump is dropped while the JS ReadableStream source is suspended awaiting more
+  // data (e.g. the client disconnected and the HTTP layer dropped the response-body pump),
+  // the underlying source's cancel() algorithm must still run. Before the fix, dropping the
+  // pump coroutine only released the reader lock and the JS source ran to natural completion.
+
+  struct TestSink final: public WritableStreamSink {
+    kj::Own<kj::PromiseFulfiller<void>> gotFirstWrite;
+    bool fired = false;
+    TestSink(kj::Own<kj::PromiseFulfiller<void>> gotFirstWrite)
+        : gotFirstWrite(kj::mv(gotFirstWrite)) {}
+
+    void signal() {
+      if (!fired) {
+        fired = true;
+        gotFirstWrite->fulfill();
+      }
+    }
+    kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override {
+      signal();
+      return kj::READY_NOW;
+    }
+    kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) override {
+      signal();
+      return kj::READY_NOW;
+    }
+    kj::Promise<void> end() override {
+      return kj::READY_NOW;
+    }
+    void abort(kj::Exception reason) override {}
+  };
+
+  capnp::MallocMessageBuilder flagsBuilder;
+  auto featureFlags = flagsBuilder.initRoot<CompatibilityFlags>();
+  featureFlags.setStreamsJavaScriptControllers(true);
+  TestFixture testFixture({.featureFlags = featureFlags.asReader()});
+
+  // Declared outside runInIoContext because the JS cancel() callback captures them by
+  // reference and runs asynchronously, after the pump is dropped.
+  bool cancelCalled = false;
+  auto cancelObserved = kj::newPromiseAndFulfiller<void>();
+
+  testFixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise<void> {
+    auto& js = jsg::Lock::from(env.isolate);
+
+    auto firstWrite = kj::newPromiseAndFulfiller<void>();
+
+    auto stream = ReadableStream::constructor(js,
+        UnderlyingSource{
+          .start =
+              [](jsg::Lock& js, auto controller) {
+      auto& c = KJ_REQUIRE_NONNULL(
+          controller.template tryGet<jsg::Ref<ReadableStreamDefaultController>>());
+      // Enqueue one chunk so the pump makes progress, but don't close the stream.
+      c->enqueue(js, jsg::JsValue(v8::ArrayBuffer::New(js.v8Isolate, 10)));
+      return js.resolvedPromise();
+    },
+          .pull =
+              [](jsg::Lock& js, auto controller) {
+      // Never enqueue more, so once the first chunk is drained the pump's next read suspends.
+      return js.resolvedPromise();
+    },
+          .cancel = [&cancelCalled, &cancelObserved](
+                        jsg::Lock& js, jsg::JsValue) -> jsg::Promise<void> {
+      cancelCalled = true;
+      if (cancelObserved.fulfiller->isWaiting()) {
+        cancelObserved.fulfiller->fulfill();
+      }
+      return js.resolvedPromise();
+    },
+        },
+        kj::none);
+
+    auto sink = kj::heap<TestSink>(kj::mv(firstWrite.fulfiller));
+    auto pump = stream->pumpTo(js, kj::mv(sink), true);
+
+    // Once the first chunk has been written the source is suspended on its next read. Drop the
+    // pump (simulating the disconnect), then wait for the source's cancel() to run.
+    return firstWrite.promise.then(
+        [pump = kj::mv(pump), cancelPromise = kj::mv(cancelObserved.promise)]() mutable {
+      { auto dropped = kj::mv(pump); }
+      return kj::mv(cancelPromise);
+    });
+  });
+
+  KJ_ASSERT(cancelCalled);
+}
+
+KJ_TEST("ReadableStream pumpTo does not cancel the JS source on clean completion") {
+  // Companion to the drop-mid-stream test above: pumpSettled must suppress the teardown
+  // cancel when the pump ran to completion before its promise was dropped.
+
+  struct NoopSink final: public WritableStreamSink {
+    kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override {
+      return kj::READY_NOW;
+    }
+    kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) override {
+      return kj::READY_NOW;
+    }
+    kj::Promise<void> end() override {
+      return kj::READY_NOW;
+    }
+    void abort(kj::Exception reason) override {}
+  };
+
+  capnp::MallocMessageBuilder flagsBuilder;
+  auto featureFlags = flagsBuilder.initRoot<CompatibilityFlags>();
+  featureFlags.setStreamsJavaScriptControllers(true);
+  TestFixture testFixture({.featureFlags = featureFlags.asReader()});
+
+  bool cancelCalled = false;
+
+  testFixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise<void> {
+    auto& js = jsg::Lock::from(env.isolate);
+
+    auto stream = ReadableStream::constructor(js,
+        UnderlyingSource{
+          .start =
+              [](jsg::Lock& js, auto controller) {
+      auto& c = KJ_REQUIRE_NONNULL(
+          controller.template tryGet<jsg::Ref<ReadableStreamDefaultController>>());
+      c->enqueue(js, jsg::JsValue(v8::ArrayBuffer::New(js.v8Isolate, 10)));
+      c->close(js);
+      return js.resolvedPromise();
+    },
+          .cancel = [&cancelCalled](jsg::Lock& js, jsg::JsValue) -> jsg::Promise<void> {
+      cancelCalled = true;
+      return js.resolvedPromise();
+    },
+        },
+        kj::none);
+
+    auto pump = stream->pumpTo(js, kj::heap<NoopSink>(), true);
+
+    // After the pump settles, turn the event loop a couple of times so an incorrectly
+    // scheduled teardown cancel would get a chance to run before asserting.
+    return env.context.waitForDeferredProxy(kj::mv(pump))
+        .then([]() {
+      return kj::evalLater([] {});
+    }).then([]() { return kj::evalLater([] {}); });
+  });
+
+  KJ_ASSERT(!cancelCalled);
+}
+
 }  // namespace
 }  // namespace workerd::api

src/workerd/api/streams/standard.c++

@@ -3356,8 +3356,8 @@ class AllReader {
 // pumped synchronously as many times as possible.
 //
 // The pump loop is a kj coroutine. Dropping the returned kj::Promise drops the
-// coroutine frame, which destroys the DrainingReader (releasing the stream lock)
-// and the sink. No WeakRef/IoOwn dance is needed because ownership is clear.
+// coroutine frame, which destroys the DrainingReader (releasing the stream lock) and the
+// sink; on a mid-stream drop the KJ_DEFER below schedules the source's cancel() first.
 // The coroutine that implements the pump loop takes ownership of the DrainingReader
 // and sink. The jsg::Ref<ReadableStream> is not passed into the coroutine because
 // jsg::Ref is disallowed in coroutine parameters; instead, the DrainingReader holds
@@ -3369,6 +3369,27 @@ kj::Promise<void> pumpToImpl(IoContext& ioContext,
 
   bool writeFailed = false;
 
+  // If the pump is dropped mid-stream (e.g. client disconnect), cancel the source so its
+  // JS cancel() algorithm runs. The isolate lock is unavailable during teardown, so the
+  // cancel is scheduled as a waitUntil task, which IncomingRequest::drain() waits for.
+  // The drop may happen during ~IoContext itself (pumps can be owned by the context's
+  // task sets), hence the WeakRef guard. The clean exits below set pumpSettled to
+  // suppress the cancel for an already-settled stream.
+  bool pumpSettled = false;
+  auto contextWeakRef = ioContext.getWeakRef();
+  KJ_DEFER({
+    if (!pumpSettled && reader->isAttached()) {
+      contextWeakRef->runIfAlive([&](IoContext& context) {
+        context.addWaitUntil(
+            context.run([reader = kj::mv(reader)](jsg::Lock& js) mutable -> kj::Promise<void> {
+          auto& ioContext = IoContext::current();
+          auto promise = ioContext.awaitJs(js, reader->cancel(js, kj::none));
+          return promise.attach(kj::mv(reader));
+        }));
+      });
+    }
+  });
+
   KJ_TRY {
     while (true) {
       // Perform a draining read to get all synchronously available data if possible
@@ -3395,6 +3416,7 @@ kj::Promise<void> pumpToImpl(IoContext& ioContext,
         if (end) {
           co_await sink->end();
         }
+        pumpSettled = true;
         co_return;
       }
     }
@@ -3409,6 +3431,9 @@ kj::Promise<void> pumpToImpl(IoContext& ioContext,
       auto error = js.exceptionToJsValue(kj::mv(ex));
       return ioContext.awaitJs(js, reader->cancel(js, error.getHandle(js)));
     });
+    // Set only after the error-path cancel completes; a drop mid-co_await lets KJ_DEFER
+    // schedule a second cancel, which is benign since controller cancel is idempotent.
+    pumpSettled = true;
     kj::throwFatalException(kj::mv(exception));
   }
 }