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)); }