From 580c53c3bb69c40acef333e500627b41583df536 Mon Sep 17 00:00:00 2001 From: JT Olio Date: Fri, 18 Sep 2026 12:21:36 +0000 Subject: [PATCH] Use weak self captures for refcounted continuations --- src/workerd/api/queue.c++ | 17 +++++----- .../api/streams/identity-transform-stream.c++ | 8 +++-- src/workerd/io/actor-sqlite.c++ | 5 +-- src/workerd/io/io-channels.c++ | 16 +++++---- src/workerd/io/io-context.h | 34 +++++++++++-------- src/workerd/server/channel-token-test.c++ | 25 ++++++++------ src/workerd/server/server.c++ | 7 ++-- 7 files changed, 64 insertions(+), 48 deletions(-) diff --git a/src/workerd/api/queue.c++ b/src/workerd/api/queue.c++ index 33d4f32412d..769152f35c8 100644 --- a/src/workerd/api/queue.c++ +++ b/src/workerd/api/queue.c++ @@ -847,22 +847,23 @@ kj::Promise QueueCustomEvent::sendRpc( } } - return req.send().then([this](auto resp) { + return req.send().then([self = addWeakToThis()](auto resp) { + auto& event = self.assertLive(); auto respResult = resp.getResult(); - this->result->ackAll = respResult.getAckAll(); + event.result->ackAll = respResult.getAckAll(); auto retryBatch = respResult.getRetryBatch(); - this->result->retryBatch.retry = retryBatch.getRetry(); + event.result->retryBatch.retry = retryBatch.getRetry(); if (retryBatch.isDelaySeconds()) { - this->result->retryBatch.delaySeconds = retryBatch.getDelaySeconds(); + event.result->retryBatch.delaySeconds = retryBatch.getDelaySeconds(); } - this->result->explicitAcks.clear(); + event.result->explicitAcks.clear(); for (const auto& msgId: respResult.getExplicitAcks()) { - this->result->explicitAcks.insert(kj::heapString(msgId)); + event.result->explicitAcks.insert(kj::heapString(msgId)); } - this->result->retries.clear(); + event.result->retries.clear(); for (const auto& retry: respResult.getRetryMessages()) { - auto& entry = this->result->retries.upsert(kj::heapString(retry.getMsgId()), {}); + auto& entry = event.result->retries.upsert(kj::heapString(retry.getMsgId()), {}); if (retry.isDelaySeconds()) { entry.value.delaySeconds = retry.getDelaySeconds(); } diff --git a/src/workerd/api/streams/identity-transform-stream.c++ b/src/workerd/api/streams/identity-transform-stream.c++ index 90f91cb9d83..5cf2d40d858 100644 --- a/src/workerd/api/streams/identity-transform-stream.c++ +++ b/src/workerd/api/streams/identity-transform-stream.c++ @@ -114,16 +114,18 @@ class IdentityTransformStreamImpl final: public kj::Refcounted, auto promise = readHelper(kj::arrayPtr(static_cast(buffer), maxBytes)); KJ_IF_SOME(l, limit) { - promise = promise.then([this, &l = l](size_t amount) -> kj::Promise { + promise = + promise.then([self = addWeakToThis(), &l = l](size_t amount) -> kj::Promise { + auto& stream = self.assertLive(); if (amount > l) { auto exception = JSG_KJ_EXCEPTION( FAILED, TypeError, "Attempt to write too many bytes through a FixedLengthStream."); - cancel(exception.clone()); + stream.cancel(exception.clone()); return kj::mv(exception); } else if (amount == 0 && l != 0) { auto exception = JSG_KJ_EXCEPTION(FAILED, TypeError, "FixedLengthStream did not see all expected bytes before close()."); - cancel(exception.clone()); + stream.cancel(exception.clone()); return kj::mv(exception); } l -= amount; diff --git a/src/workerd/io/actor-sqlite.c++ b/src/workerd/io/actor-sqlite.c++ index 2d016268911..47fec41ce55 100644 --- a/src/workerd/io/actor-sqlite.c++ +++ b/src/workerd/io/actor-sqlite.c++ @@ -195,8 +195,9 @@ kj::Maybe> ActorSqlite::ExplicitTxn::commit() { // Although the promise returned here was originally intended for "backpressure", it turns out // if we return a promise here, the one call site (DurableObjectStorage::asyncTransactionImpl()) // will actually keep the input gate locked until the commit finishes, which is what we need. - return actorSqlite.blockTasks.onEmpty().then([this]() { - commitImpl(); + return actorSqlite.blockTasks.onEmpty() + .then([self = addWeakToThis()]() { + self.assertLive().commitImpl(); }).catch_([self = kj::addRef(*this)](kj::Exception&& e) mutable { if (self->actorSqlite.broken == kj::none) { self->rollbackImpl(); diff --git a/src/workerd/io/io-channels.c++ b/src/workerd/io/io-channels.c++ index eef90e36548..cdeb2da459e 100644 --- a/src/workerd/io/io-channels.c++ +++ b/src/workerd/io/io-channels.c++ @@ -85,8 +85,10 @@ class PromisedTokenizableChannel: public ChannelType { KJ_IF_SOME(channel, inner) { return channel->getTokenMaybeSync(usage); } else { - return readyPromise.addBranch().then([this, usage]() -> kj::Promise> { - KJ_SWITCH_ONEOF(KJ_ASSERT_NONNULL(inner)->getTokenMaybeSync(usage)) { + return readyPromise.addBranch().then( + [self = this->addWeakToThis(), usage]() -> kj::Promise> { + auto& channel = self.assertLive(); + KJ_SWITCH_ONEOF(KJ_ASSERT_NONNULL(channel.inner)->getTokenMaybeSync(usage)) { KJ_CASE_ONEOF(token, kj::Array) { return kj::mv(token); } @@ -105,8 +107,9 @@ class PromisedTokenizableChannel: public ChannelType { KJ_IF_SOME(channel, inner) { return kj::addRef(*channel); } else { - return readyPromise.addBranch().then([this]() mutable { - return kj::addRef(*KJ_ASSERT_NONNULL(inner)); + return readyPromise.addBranch().then([self = this->addWeakToThis()]() mutable { + auto& channel = self.assertLive(); + return kj::addRef(*KJ_ASSERT_NONNULL(channel.inner)); }); } } @@ -166,8 +169,9 @@ class PromisedRpcChannel final: public PromisedTokenizableChannelrestore(); } else { auto splitPromise = readyPromise.addBranch() - .then([this]() { - auto innerRestore = KJ_ASSERT_NONNULL(inner)->restore(); + .then([self = addWeakToThis()]() { + auto& channel = self.assertLive(); + auto innerRestore = KJ_ASSERT_NONNULL(channel.inner)->restore(); return kj::tuple(kj::mv(innerRestore.cap), kj::mv(innerRestore.task)); }).split(); return { diff --git a/src/workerd/io/io-context.h b/src/workerd/io/io-context.h index 3be9e89ca28..89da7f0d9ce 100644 --- a/src/workerd/io/io-context.h +++ b/src/workerd/io/io-context.h @@ -842,8 +842,8 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler // into a regular `Promise`, including registering pending events as needed. template kj::Promise waitForDeferredProxy(kj::Promise>&& promise) { - return promise.then([this](api::DeferredProxy deferredProxy) { - return deferredProxy.proxyTask.attach(registerPendingEvent()); + return promise.then([self = addWeakToThis()](api::DeferredProxy deferredProxy) { + return deferredProxy.proxyTask.attach(self.assertLive().registerPendingEvent()); }); } @@ -1384,8 +1384,9 @@ kj::PromiseForResult IoContext::runSingle( KJ_IF_SOME(cs, criticalSection) { return cs.get() ->wait(getCurrentTraceSpan()) - .then([this, func = kj::fwd(func)](InputGate::Lock&& inputLock) mutable { - return runSingle(kj::fwd(func), kj::mv(inputLock)); + .then([self = addWeakToThis(), func = kj::fwd(func)]( + InputGate::Lock&& inputLock) mutable { + return self.assertLive().runSingle(kj::fwd(func), kj::mv(inputLock)); }); } else { return runSingle(kj::fwd(func)); @@ -1406,8 +1407,9 @@ kj::PromiseForResult IoContext::runSingle( if (inputLock == kj::none) { return a.getInputGate() .wait(getCurrentTraceSpan()) - .then([this, func = kj::fwd(func)](InputGate::Lock&& inputLock) mutable { - return runSingle(kj::fwd(func), kj::mv(inputLock)); + .then([self = addWeakToThis(), func = kj::fwd(func)]( + InputGate::Lock&& inputLock) mutable { + return self.assertLive().runSingle(kj::fwd(func), kj::mv(inputLock)); }); } @@ -1416,10 +1418,11 @@ kj::PromiseForResult IoContext::runSingle( asyncLockPromise = worker->takeAsyncLock(getMetrics()); } - return asyncLockPromise.then([this, inputLock = kj::mv(inputLock), func = kj::fwd(func)]( - Worker::AsyncLock lock) mutable { + return asyncLockPromise.then([self = addWeakToThis(), inputLock = kj::mv(inputLock), + func = kj::fwd(func)](Worker::AsyncLock lock) mutable { + auto& context = self.assertLive(); // Re-check if context was aborted while we waited for the lock. - KJ_IF_SOME(ex, abortException) { + KJ_IF_SOME(ex, context.abortException) { kj::throwFatalException(ex.clone()); } @@ -1436,7 +1439,7 @@ kj::PromiseForResult IoContext::runSingle( }; RunnableImpl runnable(kj::fwd(func)); - runImpl(runnable, lock, kj::mv(inputLock), Runnable::Exceptional(false)); + context.runImpl(runnable, lock, kj::mv(inputLock), Runnable::Exceptional(false)); } else { struct RunnableImpl: public Runnable { Func func; @@ -1449,7 +1452,7 @@ kj::PromiseForResult IoContext::runSingle( }; RunnableImpl runnable{kj::fwd(func)}; - runImpl(runnable, lock, kj::mv(inputLock), Runnable::Exceptional(false)); + context.runImpl(runnable, lock, kj::mv(inputLock), Runnable::Exceptional(false)); KJ_IF_SOME(r, runnable.result) { return kj::mv(r); } else { @@ -1890,18 +1893,19 @@ jsg::PromiseForResult IoContext::blockConcurrencyWhileImpl( maybeAsyncContext = jsg::AsyncContextFrame::currentRef(js)]( InputGate::Lock inputLock) mutable { return run( - [this, callback = kj::mv(callback), maybeAsyncContext = kj::mv(maybeAsyncContext)]( - Worker::Lock& lock) mutable { + [self = addWeakToThis(), callback = kj::mv(callback), + maybeAsyncContext = kj::mv(maybeAsyncContext)](Worker::Lock& lock) mutable { jsg::AsyncContextFrame::Scope scope(lock, maybeAsyncContext); auto cb = kj::mv(callback); // Remember that this can throw synchronously, and it's important that we catch such throws // and call cs->failed(). auto promise = cb(lock); + auto& context = self.assertLive(); // Arrange to time out if the critical section runs more than 30 seconds, so that objects // won't be hung forever if they have a critical section that deadlocks. - auto timeout = afterLimitTimeout(30 * kj::SECONDS).then([]() -> T { + auto timeout = context.afterLimitTimeout(30 * kj::SECONDS).then([]() -> T { auto e = JSG_KJ_EXCEPTION(OVERLOADED, Error, "A call to blockConcurrencyWhile() in a Durable Object waited for " "too long. The call was canceled and the Durable Object was reset."); @@ -1909,7 +1913,7 @@ jsg::PromiseForResult IoContext::blockConcurrencyWhileImpl( kj::throwFatalException(kj::mv(e)); }); - return awaitJs(lock, kj::mv(promise)).exclusiveJoin(kj::mv(timeout)); + return context.awaitJs(lock, kj::mv(promise)).exclusiveJoin(kj::mv(timeout)); }, kj::mv(inputLock)); }) diff --git a/src/workerd/server/channel-token-test.c++ b/src/workerd/server/channel-token-test.c++ index cda9f991080..c363caa22fb 100644 --- a/src/workerd/server/channel-token-test.c++ +++ b/src/workerd/server/channel-token-test.c++ @@ -81,10 +81,11 @@ class MockSubrequestChannel: public IoChannelFactory::SubrequestChannel { KJ_IF_SOME(p, readyPromise) { auto promise = kj::mv(p); readyPromise = kj::none; - return promise.then([&h, usage, this]() mutable -> kj::Array { - return expectSync(h.encodeSubrequestChannelToken(usage, triplet.serviceName, - triplet.entrypoint.map([](kj::String& s) -> kj::StringPtr { return s; }), triplet.props, - persistent)); + return promise.then([&h, usage, self = addWeakToThis()]() mutable -> kj::Array { + auto& channel = self.assertLive(); + return expectSync(h.encodeSubrequestChannelToken(usage, channel.triplet.serviceName, + channel.triplet.entrypoint.map([](kj::String& s) -> kj::StringPtr { return s; }), + channel.triplet.props, channel.persistent)); }); } else { return expectSync(h.encodeSubrequestChannelToken(usage, triplet.serviceName, @@ -123,10 +124,11 @@ class MockActorClassChannel: public IoChannelFactory::ActorClassChannel { KJ_IF_SOME(p, readyPromise) { auto promise = kj::mv(p); readyPromise = kj::none; - return promise.then([&h, usage, this]() mutable -> kj::Array { - return expectSync(h.encodeActorClassChannelToken(usage, triplet.serviceName, - triplet.entrypoint.map([](kj::String& s) -> kj::StringPtr { return s; }), triplet.props, - persistent)); + return promise.then([&h, usage, self = addWeakToThis()]() mutable -> kj::Array { + auto& channel = self.assertLive(); + return expectSync(h.encodeActorClassChannelToken(usage, channel.triplet.serviceName, + channel.triplet.entrypoint.map([](kj::String& s) -> kj::StringPtr { return s; }), + channel.triplet.props, channel.persistent)); }); } else { return expectSync(h.encodeActorClassChannelToken(usage, triplet.serviceName, @@ -179,9 +181,10 @@ class MockActorChannel: public IoChannelFactory::ActorChannel { KJ_IF_SOME(p, readyPromise) { auto promise = kj::mv(p); readyPromise = kj::none; - return promise.then([&h, usage, this]() mutable -> kj::Array { - return h.encodeActorChannelToken(usage, namespaceKey, id, - name.map([](kj::String& s) -> kj::StringPtr { return s; }), persistent); + return promise.then([&h, usage, self = addWeakToThis()]() mutable -> kj::Array { + auto& channel = self.assertLive(); + return h.encodeActorChannelToken(usage, channel.namespaceKey, channel.id, + channel.name.map([](kj::String& s) -> kj::StringPtr { return s; }), channel.persistent); }); } else { return h.encodeActorChannelToken(usage, namespaceKey, id, diff --git a/src/workerd/server/server.c++ b/src/workerd/server/server.c++ index 9da50a9891c..2b4a34f2feb 100644 --- a/src/workerd/server/server.c++ +++ b/src/workerd/server/server.c++ @@ -811,9 +811,10 @@ class Server::ActorNamespace final { return getChannelTokenImpl(usage, c.id, persistent); } KJ_CASE_ONEOF(promise, kj::ForkedPromise) { - return promise.addBranch().then([this, usage, persistent]() { - return getChannelTokenImpl( - usage, KJ_ASSERT_NONNULL(classAndId.tryGet()).id, persistent); + return promise.addBranch().then([self = addWeakToThis(), usage, persistent]() { + auto& container = self.assertLive(); + return container.getChannelTokenImpl( + usage, KJ_ASSERT_NONNULL(container.classAndId.tryGet()).id, persistent); }); } }