From 8fe8fa68f0b93cf8ae90f20c4ac4b975158b5a8c 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; the error path claims the suppression before awaiting its own cancel, so a drop while that is pending cannot start a second one. Both cancel sites tolerate a rejecting cancel: cancelling a stream whose source already errored rejects with the stored error, so the error path ignores that rejection to let the failure that actually broke the pump propagate, and the teardown task swallows it so best-effort disconnect cleanup cannot land as a failed waitUntil task when the stream errored between the drop and the task acquiring the isolate lock. Cancelling with kj::none (undefined reason) matches the pre-regression PumpToReader drop behavior. Regression tests pin: the drop-mid-stream cancel running, clean completion not invoking cancel(), the original sink failure surviving a rejecting cancel algorithm, the already-errored teardown window leaving the waitUntil status untouched, and request drain completing only after the teardown cancel has finished. Fixes #6832 Co-Authored-By: Claude Fable 5 --- docs/streams.md | 10 +- src/workerd/api/streams-test.c++ | 354 +++++++++++++++++++++++++++ src/workerd/api/streams/standard.c++ | 39 ++- 3 files changed, 397 insertions(+), 6 deletions(-) diff --git a/docs/streams.md b/docs/streams.md index 54f60e7e64c..7f589ea7360 100644 --- a/docs/streams.md +++ b/docs/streams.md @@ -88,7 +88,10 @@ 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 the runtime discards its own consumer of the stream (e.g. workerd observes the client + disconnect and drops the response-body pump; a JS reader merely releasing its lock does not + cancel) - **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 +416,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..ebce7b9d271 100644 --- a/src/workerd/api/streams-test.c++ +++ b/src/workerd/api/streams-test.c++ @@ -198,5 +198,359 @@ 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); +} + +KJ_TEST("ReadableStream pumpTo propagates the original failure when cancel also rejects") { + // The error path cancels the source before rethrowing. A rejecting cancel algorithm must not + // replace the failure that actually broke the pump. + + struct FailingSink final: public WritableStreamSink { + kj::Promise write(kj::ArrayPtr buffer) override { + return JSG_KJ_EXCEPTION(FAILED, Error, "sink failure"); + } + kj::Promise write(kj::ArrayPtr> pieces) override { + return JSG_KJ_EXCEPTION(FAILED, Error, "sink failure"); + } + 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 continuation outlives the lambda frame. + kj::Maybe failure; + 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))); + return js.resolvedPromise(); + }, + .cancel = [&cancelCalled](jsg::Lock& js, jsg::JsValue) -> jsg::Promise { + cancelCalled = true; + return js.rejectedPromise(js.typeError("cancel failure"_kj)); + }, + }, + kj::none); + + auto pump = stream->pumpTo(js, kj::heap(), true); + return env.context.waitForDeferredProxy(kj::mv(pump)).then([]() { + KJ_FAIL_ASSERT("pump should have failed"); + }, [&failure](kj::Exception&& e) { failure = kj::str(e.getDescription()); }); + }); + + auto& description = KJ_ASSERT_NONNULL(failure); + KJ_ASSERT(description.contains("sink failure"), description); + KJ_ASSERT(cancelCalled); +} + } // namespace +KJ_TEST("ReadableStream pumpTo teardown cancel tolerates an already-errored stream") { + // A client can disconnect right after the source errored the stream, while the pump is still + // suspended on a read that has not yet observed the error. The teardown defer then cancels a + // stream that is already errored, and that cancel rejects with the stored error. Cancellation + // is best-effort cleanup: the rejection must not land as a failed waitUntil task, which would + // report the disconnect as a request failure. + + 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 {} + }; + + // Own the IO loop so the isolate lock can be released after the pump is dropped: the teardown + // cancel only runs once the pump's own lock is gone. + auto io = kj::setupAsyncIo(); + capnp::MallocMessageBuilder flagsBuilder; + auto featureFlags = flagsBuilder.initRoot(); + featureFlags.setStreamsJavaScriptControllers(true); + TestFixture testFixture({.waitScope = io.waitScope, .featureFlags = featureFlags.asReader()}); + + auto context = testFixture.newIoContext(); + auto request = testFixture.newIncomingRequest(*context); + + context + ->run([&](Worker::Lock& lock) { + auto& js = jsg::Lock::from(lock.getIsolate()); + + kj::Maybe> savedController; + auto stream = ReadableStream::constructor(js, + UnderlyingSource{ + .start = + [&savedController](jsg::Lock& js, auto controller) { + auto& c = KJ_REQUIRE_NONNULL( + controller.template tryGet>()); + savedController = c.addRef(); + return js.resolvedPromise(); + }, + .pull = + [](jsg::Lock& js, auto controller) { + // Never enqueue, so the pump's draining read suspends and the stream stays healthy until + // the explicit error below. + return js.resolvedPromise(); + }, + }, + StreamQueuingStrategy{.highWaterMark = 0}); + + auto pump = stream->pumpTo(js, kj::heap(), true); + + // Error the stream while the pump is suspended, then drop the pump before the suspended read + // is rejected, so the teardown defer is what cancels the (already errored) stream. + KJ_ASSERT_NONNULL(savedController)->error(js, jsg::JsValue(js.str("source failure"_kj))); + { auto dropped = kj::mv(pump); } + }).wait(io.waitScope); + + // The teardown cancel task is queued ahead of this run, so it has run (and rejected) by the + // time this does. + context->run([](Worker::Lock&) {}).wait(io.waitScope); + + KJ_ASSERT(context->waitUntilStatus() == EventOutcome::OK); +} + +KJ_TEST("ReadableStream pumpTo teardown cancel completes before request drain finishes") { + // The teardown cancel must be a waitUntil task: IncomingRequest::drain() waits only for + // waitUntil tasks, and the context (with any plain tasks still pending) goes away right after + // drain completes. A plain task would let drain finish before the cancel ran, silently + // dropping the disconnect cleanup. + + 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 {} + }; + + struct DrainErrorHandler final: public kj::TaskSet::ErrorHandler { + void taskFailed(kj::Exception&& exception) override { + KJ_FAIL_EXPECT("drain task failed", exception); + } + }; + + // Own the IO loop so the isolate lock can be released after the pump is dropped, as in the + // already-errored test above. + auto io = kj::setupAsyncIo(); + capnp::MallocMessageBuilder flagsBuilder; + auto featureFlags = flagsBuilder.initRoot(); + featureFlags.setStreamsJavaScriptControllers(true); + TestFixture testFixture({.waitScope = io.waitScope, .featureFlags = featureFlags.asReader()}); + + // The request is the context's only owner, as in production, so `context` dangles once the + // completed drain destroys the request; do not touch it after the drain. + auto request = testFixture.newIncomingRequest(); + auto& context = request->getContext(); + + bool cancelStarted = false; + bool cancelCompleted = false; + auto cancelGate = kj::newPromiseAndFulfiller(); + + context + .run([&](Worker::Lock& lock) { + auto& js = jsg::Lock::from(lock.getIsolate()); + + auto stream = ReadableStream::constructor(js, + UnderlyingSource{ + .pull = + [](jsg::Lock& js, auto controller) { + // Never enqueue, so the pump's draining read suspends and the drop below is mid-stream. + return js.resolvedPromise(); + }, + .cancel = [&cancelStarted, &cancelCompleted, promise = kj::mv(cancelGate.promise)]( + jsg::Lock& js, jsg::JsValue) mutable -> jsg::Promise { + cancelStarted = true; + // Hold the cancel open on an external gate: a fast cancel would finish during the drain + // wait below under either scheduling, hiding whether drain actually waited for it. + return IoContext::current().awaitIo( + js, promise.then([&cancelCompleted]() { cancelCompleted = true; })); + }, + }, + StreamQueuingStrategy{.highWaterMark = 0}); + + auto pump = stream->pumpTo(js, kj::heap(), true); + { auto dropped = kj::mv(pump); } + }).wait(io.waitScope); + + DrainErrorHandler errorHandler; + kj::TaskSet drainTasks(errorHandler); + request->drain(drainTasks, kj::mv(request)); + auto drained = drainTasks.onEmpty(); + + KJ_ASSERT(!drained.poll(io.waitScope)); + KJ_ASSERT(cancelStarted); + KJ_ASSERT(!cancelCompleted); + + cancelGate.fulfiller->fulfill(); + drained.wait(io.waitScope); + KJ_ASSERT(cancelCompleted); +} + } // namespace workerd::api diff --git a/src/workerd/api/streams/standard.c++ b/src/workerd/api/streams/standard.c++ index ae704ab056c..03407bfd7a1 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,30 @@ 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 exits below set pumpSettled once the stream has + // settled or its cancel is under way. The cancel is best-effort: it rejects if the stream + // errored before the drop, and a rejected waitUntil task would report disconnect cleanup as a + // request failure. + 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)); + }).catch_([](kj::Exception&&) {})); + }); + } + }); + KJ_TRY { while (true) { // Perform a draining read to get all synchronously available data if possible @@ -3396,6 +3420,7 @@ kj::Promise pumpToImpl(IoContext& ioContext, if (end) { co_await sink->end(); } + pumpSettled = true; co_return; } } @@ -3405,10 +3430,16 @@ kj::Promise pumpToImpl(IoContext& ioContext, sink->abort(exception.clone()); } - co_await ioContext.run([&reader, ex = exception.clone()](jsg::Lock& js) mutable { + co_await ioContext.run([&reader, &pumpSettled, ex = exception.clone()](jsg::Lock& js) mutable { auto& ioContext = IoContext::current(); + // A drop while the cancel below is still pending must not make the teardown defer start a + // second one, so claim the suppression now rather than after the await. + pumpSettled = true; auto error = js.exceptionToJsValue(kj::mv(ex)); - return ioContext.awaitJs(js, reader->cancel(js, error.getHandle(js))); + // Cancelling a stream the source already errored rejects with the stored error. Ignore it + // so `exception` propagates below. + return ioContext.awaitJs(js, reader->cancel(js, error.getHandle(js))) + .catch_([](kj::Exception&&) {}); }); kj::throwFatalException(kj::mv(exception)); }