From 752d9e6d4ef97b7a8b1b90999554ef1991aa1710 Mon Sep 17 00:00:00 2001 From: Michael Hart Date: Thu, 23 Jul 2026 20:54:47 +1000 Subject: [PATCH] propagate cancellation to standard stream sources --- src/workerd/api/streams/standard-test.c++ | 59 ++++++++++++++++++++ src/workerd/api/streams/standard.c++ | 67 +++++++++++++++++------ 2 files changed, 108 insertions(+), 18 deletions(-) diff --git a/src/workerd/api/streams/standard-test.c++ b/src/workerd/api/streams/standard-test.c++ index 7360b919e62..6b334b29992 100644 --- a/src/workerd/api/streams/standard-test.c++ +++ b/src/workerd/api/streams/standard-test.c++ @@ -15,6 +15,65 @@ void preamble(auto callback) { fixture.runInIoContext([&](const TestFixture::Environment& env) { callback(env.js); }); } +class PumpTestSink final: public WritableStreamSink { + public: + kj::Promise write(kj::ArrayPtr) override { + return kj::READY_NOW; + } + + kj::Promise write(kj::ArrayPtr>) override { + return kj::READY_NOW; + } + + kj::Promise end() override { + return kj::READY_NOW; + } + + void abort(kj::Exception) override {} +}; + +KJ_TEST("standard stream pump cancellation invokes source cancel") { + TestFixture fixture; + bool cancelCalled = false; + bool pullCalled = false; + kj::Maybe::Resolver> pendingResolver; + + fixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise { + auto cancelDone = env.js.newPromiseAndResolver(); + auto rs = env.js.alloc(newReadableStreamJsController()); + rs->getController().setup(env.js, + UnderlyingSource{ + .pull = + [&pullCalled, &pendingResolver](jsg::Lock& js, UnderlyingSource::Controller) { + pullCalled = true; + auto pending = js.newPromiseAndResolver(); + pendingResolver = kj::mv(pending.resolver); + return kj::mv(pending.promise); + }, + .cancel = + [resolver = kj::mv(cancelDone.resolver), &cancelCalled](jsg::Lock& js, auto) mutable { + cancelCalled = true; + resolver.resolve(js); + return js.resolvedPromise(); + }, + }, + StreamQueuingStrategy{.highWaterMark = 1}); + KJ_ASSERT(pullCalled, "pull() should be pending before the pump starts"); + KJ_ASSERT(pendingResolver != kj::none, "pull() should leave a promise pending"); + + auto canceler = kj::heap(); + auto pump = canceler->wrap(rs->getController().pumpTo(env.js, kj::heap(), true)) + .then([](DeferredProxy proxy) { return kj::mv(proxy.proxyTask); }, + [](kj::Exception&&) {}); + canceler->cancel(KJ_EXCEPTION(DISCONNECTED, "test canceled")); + auto waitForCancel = env.context.awaitJs(env.js, kj::mv(cancelDone.promise)); + return kj::joinPromises(kj::arr(kj::mv(pump), kj::mv(waitForCancel))).then([&] { + KJ_ASSERT(pullCalled, "pull() should be called before the pump is canceled"); + KJ_ASSERT(cancelCalled, "cancel() should be called when the pump is canceled"); + }); + }); +} + jsg::JsValue toBytes(jsg::Lock& js, kj::String str) { return jsg::JsValue( jsg::BackingStore::from(js, str.asBytes().attach(kj::mv(str))).createHandle(js)); diff --git a/src/workerd/api/streams/standard.c++ b/src/workerd/api/streams/standard.c++ index ae704ab056c..40d3bb32754 100644 --- a/src/workerd/api/streams/standard.c++ +++ b/src/workerd/api/streams/standard.c++ @@ -3358,15 +3358,41 @@ class AllReader { // // 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. -// 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 -// a reference to the stream internally. -kj::Promise pumpToImpl(IoContext& ioContext, - kj::Own reader, - kj::Own sink, - bool end) { +// and the sink. The cleanup state also schedules a JavaScript-level cancellation +// while the IoContext is still available, so a pending pull is not left running. +// KJ promises are event-loop-affine, so PumpState destruction occurs on the +// IoContext's event-loop thread. +class PumpState final { + public: + PumpState(IoContext& ioContext, kj::Own reader, kj::Own sink) + : ioContext(ioContext.getWeakRef()), + reader(ioContext.addObject(kj::mv(reader))), + sink(kj::mv(sink)) {} + + ~PumpState() noexcept(false) { + if (done || cleanupStarted) return; + + auto reader = kj::mv(this->reader); + ioContext->runIfAlive([reader = kj::mv(reader)](IoContext& ioContext) mutable { + if (!ioContext.hasCurrentIncomingRequest()) return; + ioContext.addTask( + ioContext.run([reader = kj::mv(reader)](jsg::Lock& js, IoContext& ioContext) mutable { + auto reason = js.error("ReadableStream pump canceled"); + return ioContext.awaitJs(js, reader->cancel(js, reason)).catch_([](kj::Exception&&) {}); + })); + }); + } + + KJ_DISALLOW_COPY_AND_MOVE(PumpState); + + kj::Own ioContext; + IoOwn reader; + kj::Own sink; + bool done = false; + bool cleanupStarted = false; +}; + +kj::Promise pumpToImpl(IoContext& ioContext, kj::Own state, bool end) { bool writeFailed = false; @@ -3374,12 +3400,12 @@ kj::Promise pumpToImpl(IoContext& ioContext, while (true) { // Perform a draining read to get all synchronously available data if possible // or fall back to a regular read if not. - DrainingReadResult result = co_await ioContext.run([&reader](jsg::Lock& js) mutable { + DrainingReadResult result = co_await ioContext.run([&state](jsg::Lock& js) mutable { auto& ioContext = IoContext::current(); // Use a 256KB limit to allow periodic yielding to the event loop, // preventing a fast producer from monopolizing the thread. constexpr size_t kMaxReadPerCycle = 256 * 1024; - return ioContext.awaitJs(js, reader->read(js, kMaxReadPerCycle)); + return ioContext.awaitJs(js, state->reader->read(js, kMaxReadPerCycle)); }); // Write all the chunks we received using vectored write for efficiency. @@ -3387,29 +3413,33 @@ kj::Promise pumpToImpl(IoContext& ioContext, KJ_ON_SCOPE_FAILURE(writeFailed = true); auto pieces = KJ_MAP(chunk, result.chunks) -> kj::ArrayPtr { return chunk.asPtr(); }; - co_await sink->write(pieces); + co_await state->sink->write(pieces); } // If the stream is done, end the output if needed and exit. if (result.done) { KJ_ON_SCOPE_FAILURE(writeFailed = true); if (end) { - co_await sink->end(); + co_await state->sink->end(); } + state->done = true; co_return; } } } KJ_CATCH(exception) { if (!writeFailed) { - sink->abort(exception.clone()); + state->sink->abort(exception.clone()); } - co_await ioContext.run([&reader, ex = exception.clone()](jsg::Lock& js) mutable { + co_await ioContext + .run([&state, ex = exception.clone()](jsg::Lock& js) mutable { auto& ioContext = IoContext::current(); + state->cleanupStarted = true; auto error = js.exceptionToJsValue(kj::mv(ex)); - return ioContext.awaitJs(js, reader->cancel(js, error.getHandle(js))); - }); + return ioContext.awaitJs(js, state->reader->cancel(js, error.getHandle(js))); + }).catch_([](kj::Exception&&) {}); + state->done = true; kj::throwFatalException(kj::mv(exception)); } } @@ -3559,7 +3589,8 @@ kj::Promise> ReadableStreamJsController::pumpTo( auto reader = KJ_ASSERT_NONNULL(DrainingReader::create(js, *this->addRef()), "Failed to create DrainingReader — stream should not be locked"); auto& ioContext = IoContext::current(); - return addNoopDeferredProxy(pumpToImpl(ioContext, kj::mv(reader), kj::mv(sink), end)); + return addNoopDeferredProxy( + pumpToImpl(ioContext, kj::heap(ioContext, kj::mv(reader), kj::mv(sink)), end)); }; KJ_SWITCH_ONEOF(state) {