Refactor waitUntilTasks from CustomEvent::sendRpc. (#2865)

Using waitUntilTasks in sendRpc is a footgun as you might depend on
a request to the rpc server that will happen complete after the request
has been responded to and the client closes the connection causing a
waitUntilTasks is not empty error. This PR and the internal PR fix any
remaining usages of waitUntilTasks and finally refactors waitUntilTasks
out of the code.

src/workerd/api/hibernatable-web-socket.c++

@@ -144,7 +144,6 @@ kj::Promise<WorkerInterface::CustomEvent::Result> HibernatableWebSocketCustomEve
 kj::Promise<WorkerInterface::CustomEvent::Result> HibernatableWebSocketCustomEventImpl::sendRpc(
     capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
     capnp::ByteStreamFactory& byteStreamFactory,
-    kj::TaskSet& waitUntilTasks,
     rpc::EventDispatcher::Client dispatcher) {
   auto req = dispatcher.castAs<rpc::HibernatableWebSocketEventDispatcher>()
                  .hibernatableWebSocketEventRequest();
@@ -194,15 +193,12 @@ HibernatableWebSocketEvent::ItemsForRelease::ItemsForRelease(
       tags(kj::mv(tags)) {}
 
 HibernatableWebSocketCustomEventImpl::HibernatableWebSocketCustomEventImpl(uint16_t typeId,
-    kj::TaskSet& waitUntilTasks,
     kj::Own<HibernationReader> params,
     kj::Maybe<Worker::Actor::HibernationManager&> manager)
     : typeId(typeId),
       params(kj::mv(params)) {}
-HibernatableWebSocketCustomEventImpl::HibernatableWebSocketCustomEventImpl(uint16_t typeId,
-    kj::TaskSet& waitUntilTasks,
-    HibernatableSocketParams params,
-    Worker::Actor::HibernationManager& manager)
+HibernatableWebSocketCustomEventImpl::HibernatableWebSocketCustomEventImpl(
+    uint16_t typeId, HibernatableSocketParams params, Worker::Actor::HibernationManager& manager)
     : typeId(typeId),
       params(kj::mv(params)),
       manager(manager) {}

src/workerd/api/hibernatable-web-socket.h

@@ -59,21 +59,17 @@ class HibernatableWebSocketCustomEventImpl final: public WorkerInterface::Custom
                                                   public kj::Refcounted {
 public:
   HibernatableWebSocketCustomEventImpl(uint16_t typeId,
-      kj::TaskSet& waitUntilTasks,
       kj::Own<HibernationReader> params,
       kj::Maybe<Worker::Actor::HibernationManager&> manager = kj::none);
-  HibernatableWebSocketCustomEventImpl(uint16_t typeId,
-      kj::TaskSet& waitUntilTasks,
-      HibernatableSocketParams params,
-      Worker::Actor::HibernationManager& manager);
+  HibernatableWebSocketCustomEventImpl(
+      uint16_t typeId, HibernatableSocketParams params, Worker::Actor::HibernationManager& manager);
 
   kj::Promise<Result> run(kj::Own<IoContext_IncomingRequest> incomingRequest,
       kj::Maybe<kj::StringPtr> entrypointName,
       kj::TaskSet& waitUntilTasks) override;
 
   kj::Promise<Result> sendRpc(capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
       capnp::ByteStreamFactory& byteStreamFactory,
-      kj::TaskSet& waitUntilTasks,
       rpc::EventDispatcher::Client dispatcher) override;
 
   uint16_t getType() override {

src/workerd/api/queue.c++

@@ -596,7 +596,6 @@ kj::Promise<WorkerInterface::CustomEvent::Result> QueueCustomEventImpl::run(
 kj::Promise<WorkerInterface::CustomEvent::Result> QueueCustomEventImpl::sendRpc(
     capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
     capnp::ByteStreamFactory& byteStreamFactory,
-    kj::TaskSet& waitUntilTasks,
     rpc::EventDispatcher::Client dispatcher) {
   auto req = dispatcher.castAs<rpc::EventDispatcher>().queueRequest();
   KJ_SWITCH_ONEOF(params) {

src/workerd/api/queue.h

@@ -345,7 +345,6 @@ class QueueCustomEventImpl final: public WorkerInterface::CustomEvent, public kj
 
   kj::Promise<Result> sendRpc(capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
       capnp::ByteStreamFactory& byteStreamFactory,
-      kj::TaskSet& waitUntilTasks,
       rpc::EventDispatcher::Client dispatcher) override;
 
   static const uint16_t EVENT_TYPE = 5;

src/workerd/api/trace.c++

@@ -646,7 +646,6 @@ auto TraceCustomEventImpl::run(kj::Own<IoContext::IncomingRequest> incomingReque
 
 auto TraceCustomEventImpl::sendRpc(capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
     capnp::ByteStreamFactory& byteStreamFactory,
-    kj::TaskSet& waitUntilTasks,
     workerd::rpc::EventDispatcher::Client dispatcher) -> kj::Promise<Result> {
   auto req = dispatcher.sendTracesRequest();
   auto out = req.initTraces(traces.size());

src/workerd/api/trace.h

@@ -604,8 +604,7 @@ class UnsafeTraceMetrics final: public jsg::Object {
 
 class TraceCustomEventImpl final: public WorkerInterface::CustomEvent {
 public:
-  TraceCustomEventImpl(
-      uint16_t typeId, kj::TaskSet& waitUntilTasks, kj::Array<kj::Own<Trace>> traces)
+  TraceCustomEventImpl(uint16_t typeId, kj::Array<kj::Own<Trace>> traces)
       : typeId(typeId),
         traces(kj::mv(traces)) {}
 
@@ -615,7 +614,6 @@ class TraceCustomEventImpl final: public WorkerInterface::CustomEvent {
 
   kj::Promise<Result> sendRpc(capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
       capnp::ByteStreamFactory& byteStreamFactory,
-      kj::TaskSet& waitUntilTasks,
       rpc::EventDispatcher::Client dispatcher) override;
 
   uint16_t getType() override {

src/workerd/api/worker-rpc.c++

@@ -1745,7 +1745,6 @@ kj::Promise<WorkerInterface::CustomEvent::Result> JsRpcSessionCustomEventImpl::r
 kj::Promise<WorkerInterface::CustomEvent::Result> JsRpcSessionCustomEventImpl::sendRpc(
     capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
     capnp::ByteStreamFactory& byteStreamFactory,
-    kj::TaskSet& waitUntilTasks,
     rpc::EventDispatcher::Client dispatcher) {
   // We arrange to revoke all capabilities in this session as soon as `sendRpc()` completes or is
   // canceled. Normally, the server side doesn't return if any capabilities still exist, so this

src/workerd/api/worker-rpc.h

@@ -420,7 +420,6 @@ class JsRpcSessionCustomEventImpl final: public WorkerInterface::CustomEvent {
 
   kj::Promise<Result> sendRpc(capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
       capnp::ByteStreamFactory& byteStreamFactory,
-      kj::TaskSet& waitUntilTasks,
       rpc::EventDispatcher::Client dispatcher) override;
 
   uint16_t getType() override {

src/workerd/io/hibernation-manager.c++

@@ -261,7 +261,7 @@ kj::Promise<void> HibernationManagerImpl::handleSocketTermination(
     auto workerInterface = loopback->getWorker(IoChannelFactory::SubrequestMetadata{});
     event = workerInterface
                 ->customEvent(kj::heap<api::HibernatableWebSocketCustomEventImpl>(
-                    hibernationEventType, readLoopTasks, kj::mv(KJ_REQUIRE_NONNULL(params)), *this))
+                    hibernationEventType, kj::mv(KJ_REQUIRE_NONNULL(params)), *this))
                 .ignoreResult()
                 .attach(kj::mv(workerInterface));
   }
@@ -366,7 +366,7 @@ kj::Promise<void> HibernationManagerImpl::readLoop(HibernatableWebSocket& hib) {
     // Dispatch the event.
     auto workerInterface = loopback->getWorker(IoChannelFactory::SubrequestMetadata{});
     co_await workerInterface->customEvent(kj::heap<api::HibernatableWebSocketCustomEventImpl>(
-        hibernationEventType, readLoopTasks, kj::mv(params), *this));
+        hibernationEventType, kj::mv(params), *this));
     if (isClose) {
       co_return;
     }

src/workerd/io/worker-interface.c++

@@ -345,11 +345,9 @@ kj::Own<WorkerInterface> WorkerInterface::fromException(kj::Exception&& e) {
 
 RpcWorkerInterface::RpcWorkerInterface(capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
     capnp::ByteStreamFactory& byteStreamFactory,
-    kj::TaskSet& waitUntilTasks,
     rpc::EventDispatcher::Client dispatcher)
     : httpOverCapnpFactory(httpOverCapnpFactory),
       byteStreamFactory(byteStreamFactory),
-      waitUntilTasks(waitUntilTasks),
       dispatcher(kj::mv(dispatcher)) {}
 
 kj::Promise<void> RpcWorkerInterface::request(kj::HttpMethod method,
@@ -405,8 +403,7 @@ kj::Promise<WorkerInterface::AlarmResult> RpcWorkerInterface::runAlarm(
 
 kj::Promise<WorkerInterface::CustomEvent::Result> RpcWorkerInterface::customEvent(
     kj::Own<CustomEvent> event) {
-  return event->sendRpc(httpOverCapnpFactory, byteStreamFactory, waitUntilTasks, dispatcher)
-      .attach(kj::mv(event));
+  return event->sendRpc(httpOverCapnpFactory, byteStreamFactory, dispatcher).attach(kj::mv(event));
 }
 
 // ======================================================================================

src/workerd/io/worker-interface.h

@@ -116,7 +116,6 @@ class WorkerInterface: public kj::HttpService {
     // Forward the event over RPC.
     virtual kj::Promise<Result> sendRpc(capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
         capnp::ByteStreamFactory& byteStreamFactory,
-        kj::TaskSet& waitUntilTasks,
         rpc::EventDispatcher::Client dispatcher) = 0;
 
     // Get the type for this event for logging / metrics purposes. This is intended for use by the
@@ -166,7 +165,6 @@ class RpcWorkerInterface: public WorkerInterface {
 public:
   RpcWorkerInterface(capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
       capnp::ByteStreamFactory& byteStreamFactory,
-      kj::TaskSet& waitUntilTasks,
       rpc::EventDispatcher::Client dispatcher);
 
   kj::Promise<void> request(kj::HttpMethod method,
@@ -189,7 +187,6 @@ class RpcWorkerInterface: public WorkerInterface {
 private:
   capnp::HttpOverCapnpFactory& httpOverCapnpFactory;
   capnp::ByteStreamFactory& byteStreamFactory;
-  kj::TaskSet& waitUntilTasks;
   rpc::EventDispatcher::Client dispatcher;
 };
 

src/workerd/server/server.c++

@@ -705,8 +705,7 @@ private:
       auto dispatcher =
           bootstrap.startEventRequest(capnp::MessageSize{4, 0}).send().getDispatcher();
       return event
-          ->sendRpc(parent.httpOverCapnpFactory, parent.byteStreamFactory, parent.waitUntilTasks,
-              kj::mv(dispatcher))
+          ->sendRpc(parent.httpOverCapnpFactory, parent.byteStreamFactory, kj::mv(dispatcher))
           .attach(kj::mv(event));
     }