Use weak self captures for refcounted continuations

src/workerd/api/queue.c++

@@ -847,22 +847,23 @@ kj::Promise<WorkerInterface::CustomEvent::Result> 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();
       }

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<kj::byte*>(buffer), maxBytes));
 
     KJ_IF_SOME(l, limit) {
-      promise = promise.then([this, &l = l](size_t amount) -> kj::Promise<size_t> {
+      promise =
+          promise.then([self = addWeakToThis(), &l = l](size_t amount) -> kj::Promise<size_t> {
+        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;

src/workerd/io/actor-sqlite.c++

@@ -195,8 +195,9 @@ kj::Maybe<kj::Promise<void>> 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();

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::Array<byte>> {
-        KJ_SWITCH_ONEOF(KJ_ASSERT_NONNULL(inner)->getTokenMaybeSync(usage)) {
+      return readyPromise.addBranch().then(
+          [self = this->addWeakToThis(), usage]() -> kj::Promise<kj::Array<byte>> {
+        auto& channel = self.assertLive();
+        KJ_SWITCH_ONEOF(KJ_ASSERT_NONNULL(channel.inner)->getTokenMaybeSync(usage)) {
           KJ_CASE_ONEOF(token, kj::Array<byte>) {
             return kj::mv(token);
           }
@@ -105,8 +107,9 @@ class PromisedTokenizableChannel: public ChannelType {
     KJ_IF_SOME(channel, inner) {
       return kj::addRef<IoChannelFactory::TokenizableChannel>(*channel);
     } else {
-      return readyPromise.addBranch().then([this]() mutable {
-        return kj::addRef<IoChannelFactory::TokenizableChannel>(*KJ_ASSERT_NONNULL(inner));
+      return readyPromise.addBranch().then([self = this->addWeakToThis()]() mutable {
+        auto& channel = self.assertLive();
+        return kj::addRef<IoChannelFactory::TokenizableChannel>(*KJ_ASSERT_NONNULL(channel.inner));
       });
     }
   }
@@ -166,8 +169,9 @@ class PromisedRpcChannel final: public PromisedTokenizableChannel<IoChannelFacto
       return channel->restore();
     } 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 {

src/workerd/io/io-context.h

@@ -842,8 +842,8 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler
   // into a regular `Promise<T>`, including registering pending events as needed.
   template <typename T>
   kj::Promise<T> waitForDeferredProxy(kj::Promise<api::DeferredProxy<T>>&& promise) {
-    return promise.then([this](api::DeferredProxy<T> deferredProxy) {
-      return deferredProxy.proxyTask.attach(registerPendingEvent());
+    return promise.then([self = addWeakToThis()](api::DeferredProxy<T> deferredProxy) {
+      return deferredProxy.proxyTask.attach(self.assertLive().registerPendingEvent());
     });
   }
 
@@ -1384,8 +1384,9 @@ kj::PromiseForResult<Func, Worker::Lock&> IoContext::runSingle(
   KJ_IF_SOME(cs, criticalSection) {
     return cs.get()
         ->wait(getCurrentTraceSpan())
-        .then([this, func = kj::fwd<Func>(func)](InputGate::Lock&& inputLock) mutable {
-      return runSingle(kj::fwd<Func>(func), kj::mv(inputLock));
+        .then([self = addWeakToThis(), func = kj::fwd<Func>(func)](
+                  InputGate::Lock&& inputLock) mutable {
+      return self.assertLive().runSingle(kj::fwd<Func>(func), kj::mv(inputLock));
     });
   } else {
     return runSingle(kj::fwd<Func>(func));
@@ -1406,8 +1407,9 @@ kj::PromiseForResult<Func, Worker::Lock&> IoContext::runSingle(
     if (inputLock == kj::none) {
       return a.getInputGate()
           .wait(getCurrentTraceSpan())
-          .then([this, func = kj::fwd<Func>(func)](InputGate::Lock&& inputLock) mutable {
-        return runSingle(kj::fwd<Func>(func), kj::mv(inputLock));
+          .then([self = addWeakToThis(), func = kj::fwd<Func>(func)](
+                    InputGate::Lock&& inputLock) mutable {
+        return self.assertLive().runSingle(kj::fwd<Func>(func), kj::mv(inputLock));
       });
     }
 
@@ -1416,10 +1418,11 @@ kj::PromiseForResult<Func, Worker::Lock&> IoContext::runSingle(
     asyncLockPromise = worker->takeAsyncLock(getMetrics());
   }
 
-  return asyncLockPromise.then([this, inputLock = kj::mv(inputLock), func = kj::fwd<Func>(func)](
-                                   Worker::AsyncLock lock) mutable {
+  return asyncLockPromise.then([self = addWeakToThis(), inputLock = kj::mv(inputLock),
+                                   func = kj::fwd<Func>(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<Func, Worker::Lock&> IoContext::runSingle(
       };
 
       RunnableImpl runnable(kj::fwd<Func>(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<Func, Worker::Lock&> IoContext::runSingle(
       };
 
       RunnableImpl runnable{kj::fwd<Func>(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,26 +1893,27 @@ jsg::PromiseForResult<Func, void, true> 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.");
         e.setDetail(WALL_TIME_LIMIT_DETAIL_ID, kj::heapArray<kj::byte>(0));
         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));
   })

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<byte> {
-        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<byte> {
+        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<byte> {
-        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<byte> {
+        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<byte> {
-        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<byte> {
+        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,

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<void>) {
-          return promise.addBranch().then([this, usage, persistent]() {
-            return getChannelTokenImpl(
-                usage, KJ_ASSERT_NONNULL(classAndId.tryGet<ClassAndId>()).id, persistent);
+          return promise.addBranch().then([self = addWeakToThis(), usage, persistent]() {
+            auto& container = self.assertLive();
+            return container.getChannelTokenImpl(
+                usage, KJ_ASSERT_NONNULL(container.classAndId.tryGet<ClassAndId>()).id, persistent);
           });
         }
       }