From a5e76574caae4fd14226767e3b7b63d4a2e57e72 Mon Sep 17 00:00:00 2001 From: Sean Chang Date: Sat, 20 Jun 2026 18:45:52 +0800 Subject: [PATCH] Cancel the JS source when a ReadableStream pump is dropped mid-stream After the draining-read pump path replaced PumpToReader (commit c32933263), dropping the pump coroutine while the JS ReadableStream source was suspended awaiting more data no longer cancelled the source. This regressed client-disconnect handling: when the HTTP layer drops the response-body pump, the underlying source's cancel() algorithm never ran and the stream kept producing data until natural completion. Restore the behavior from pumpToImpl's teardown: when the coroutine is dropped before settling, schedule the source cancel as a waitUntil task so IncomingRequest::drain() keeps the context open until the JS cancel() runs -- the same scheduling WorkerEntrypoint uses for its disconnect abort task. The isolate lock is unavailable during coroutine teardown, and the pump promise may itself be owned by the IoContext (response-body pumps live in the context's task sets), so the defer can run during ~IoContext; an IoContext::WeakRef guard (the pattern documented at IoContext::getWeakRef()) turns that case into a safe no-op. A pumpSettled flag suppresses the cancel on the normal-completion and error paths, which already finalize the stream themselves, and a companion test pins that clean completion does not invoke cancel(). Cancelling with kj::none (undefined reason) matches the pre-regression PumpToReader drop behavior. Fixes #6832 Co-Authored-By: Claude Fable 5 --- docs/streams.md | 8 +- src/workerd/api/streams-test.c++ | 147 +++++++++++++++++++++++++++ src/workerd/api/streams/standard.c++ | 29 +++++- 3 files changed, 180 insertions(+), 4 deletions(-) diff --git a/docs/streams.md b/docs/streams.md index 54f60e7e64c..caad6965b1d 100644 --- a/docs/streams.md +++ b/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 diff --git a/src/workerd/api/streams-test.c++ b/src/workerd/api/streams-test.c++ index 257f7377613..5f1d8775b26 100644 --- a/src/workerd/api/streams-test.c++ +++ b/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> gotFirstWrite; + bool fired = false; + TestSink(kj::Own> gotFirstWrite) + : gotFirstWrite(kj::mv(gotFirstWrite)) {} + + void signal() { + if (!fired) { + fired = true; + gotFirstWrite->fulfill(); + } + } + kj::Promise write(kj::ArrayPtr buffer) override { + signal(); + return kj::READY_NOW; + } + kj::Promise write(kj::ArrayPtr> pieces) override { + signal(); + return kj::READY_NOW; + } + kj::Promise end() override { + return kj::READY_NOW; + } + void abort(kj::Exception reason) override {} + }; + + capnp::MallocMessageBuilder flagsBuilder; + auto featureFlags = flagsBuilder.initRoot(); + 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(); + + testFixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise { + auto& js = jsg::Lock::from(env.isolate); + + auto firstWrite = kj::newPromiseAndFulfiller(); + + auto stream = ReadableStream::constructor(js, + UnderlyingSource{ + .start = + [](jsg::Lock& js, auto controller) { + auto& c = KJ_REQUIRE_NONNULL( + controller.template tryGet>()); + // 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 { + cancelCalled = true; + if (cancelObserved.fulfiller->isWaiting()) { + cancelObserved.fulfiller->fulfill(); + } + return js.resolvedPromise(); + }, + }, + kj::none); + + auto sink = kj::heap(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 write(kj::ArrayPtr buffer) override { + return kj::READY_NOW; + } + kj::Promise write(kj::ArrayPtr> pieces) override { + return kj::READY_NOW; + } + kj::Promise end() override { + return kj::READY_NOW; + } + void abort(kj::Exception reason) override {} + }; + + capnp::MallocMessageBuilder flagsBuilder; + auto featureFlags = flagsBuilder.initRoot(); + featureFlags.setStreamsJavaScriptControllers(true); + TestFixture testFixture({.featureFlags = featureFlags.asReader()}); + + bool cancelCalled = false; + + testFixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise { + 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>()); + 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 { + cancelCalled = true; + return js.resolvedPromise(); + }, + }, + kj::none); + + auto pump = stream->pumpTo(js, kj::heap(), 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 diff --git a/src/workerd/api/streams/standard.c++ b/src/workerd/api/streams/standard.c++ index ae704ab056c..bf13d07763b 100644 --- a/src/workerd/api/streams/standard.c++ +++ b/src/workerd/api/streams/standard.c++ @@ -3357,8 +3357,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 is not passed into the coroutine because // jsg::Ref is disallowed in coroutine parameters; instead, the DrainingReader holds @@ -3370,6 +3370,27 @@ kj::Promise 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 { + 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 @@ -3396,6 +3417,7 @@ kj::Promise pumpToImpl(IoContext& ioContext, if (end) { co_await sink->end(); } + pumpSettled = true; co_return; } } @@ -3410,6 +3432,9 @@ kj::Promise 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)); } }