Add waitUntilTasks param to CustomEvent::run().

`JsRpcCustomEventImpl` will need this. I'm surprised it wasn't provided already.
cloudflare · Mar 27, 2024 · d127570 · d127570
1 parent 7e6f215
commit d127570
Show file tree

Hide file tree

Showing 10 changed files with 20 additions and 10 deletions.
diff --git a/src/workerd/api/hibernatable-web-socket.c++ b/src/workerd/api/hibernatable-web-socket.c++
@@ -58,7 +58,8 @@ jsg::Ref<WebSocket> HibernatableWebSocketEvent::claimWebSocket(jsg::Lock& lock,
 
 kj::Promise<WorkerInterface::CustomEvent::Result> HibernatableWebSocketCustomEventImpl::run(
     kj::Own<IoContext_IncomingRequest> incomingRequest,
-    kj::Maybe<kj::StringPtr> entrypointName) {
+    kj::Maybe<kj::StringPtr> entrypointName,
+    kj::TaskSet& waitUntilTasks) {
   // Mark the request as delivered because we're about to run some JS.
   auto& context = incomingRequest->getContext();
   incomingRequest->delivered();

diff --git a/src/workerd/api/hibernatable-web-socket.h b/src/workerd/api/hibernatable-web-socket.h
@@ -71,7 +71,8 @@ class HibernatableWebSocketCustomEventImpl final: public WorkerInterface::Custom
 
   kj::Promise<Result> run(
       kj::Own<IoContext_IncomingRequest> incomingRequest,
-      kj::Maybe<kj::StringPtr> entrypointName) override;
+      kj::Maybe<kj::StringPtr> entrypointName,
+      kj::TaskSet& waitUntilTasks) override;
 
   kj::Promise<Result> sendRpc(
       capnp::HttpOverCapnpFactory& httpOverCapnpFactory,

diff --git a/src/workerd/api/queue.c++ b/src/workerd/api/queue.c++
@@ -528,7 +528,8 @@ jsg::Ref<QueueEvent> startQueueEvent(
 
 kj::Promise<WorkerInterface::CustomEvent::Result> QueueCustomEventImpl::run(
     kj::Own<IoContext_IncomingRequest> incomingRequest,
-    kj::Maybe<kj::StringPtr> entrypointName) {
+    kj::Maybe<kj::StringPtr> entrypointName,
+    kj::TaskSet& waitUntilTasks) {
   incomingRequest->delivered();
   auto& context = incomingRequest->getContext();
 

diff --git a/src/workerd/api/queue.h b/src/workerd/api/queue.h
@@ -324,7 +324,8 @@ class QueueCustomEventImpl final: public WorkerInterface::CustomEvent, public kj
 
   kj::Promise<Result> run(
       kj::Own<IoContext_IncomingRequest> incomingRequest,
-      kj::Maybe<kj::StringPtr> entrypointName) override;
+      kj::Maybe<kj::StringPtr> entrypointName,
+      kj::TaskSet& waitUntilTasks) override;
 
   kj::Promise<Result> sendRpc(
       capnp::HttpOverCapnpFactory& httpOverCapnpFactory,

diff --git a/src/workerd/api/trace.c++ b/src/workerd/api/trace.c++
@@ -610,7 +610,8 @@ kj::Promise<void> sendTracesToExportedHandler(
 }  // namespace
 
 auto TraceCustomEventImpl::run(
-    kj::Own<IoContext::IncomingRequest> incomingRequest, kj::Maybe<kj::StringPtr> entrypointNamePtr)
+    kj::Own<IoContext::IncomingRequest> incomingRequest, kj::Maybe<kj::StringPtr> entrypointNamePtr,
+    kj::TaskSet& waitUntilTasks)
     -> kj::Promise<Result> {
   // Don't bother to wait around for the handler to run, just hand it off to the waitUntil tasks.
   waitUntilTasks.add(

diff --git a/src/workerd/api/trace.h b/src/workerd/api/trace.h
@@ -590,7 +590,8 @@ class TraceCustomEventImpl final: public WorkerInterface::CustomEvent {
 
   kj::Promise<Result> run(
       kj::Own<IoContext::IncomingRequest> incomingRequest,
-      kj::Maybe<kj::StringPtr> entrypointName) override;
+      kj::Maybe<kj::StringPtr> entrypointName,
+      kj::TaskSet& waitUntilTasks) override;
 
   kj::Promise<Result> sendRpc(
       capnp::HttpOverCapnpFactory& httpOverCapnpFactory,

diff --git a/src/workerd/api/worker-rpc.c++ b/src/workerd/api/worker-rpc.c++
@@ -1673,7 +1673,8 @@ private:
 
 kj::Promise<WorkerInterface::CustomEvent::Result> JsRpcSessionCustomEventImpl::run(
     kj::Own<IoContext::IncomingRequest> incomingRequest,
-    kj::Maybe<kj::StringPtr> entrypointName) {
+    kj::Maybe<kj::StringPtr> entrypointName,
+    kj::TaskSet& waitUntilTasks) {
   IoContext& ioctx = incomingRequest->getContext();
 
   incomingRequest->delivered();

diff --git a/src/workerd/api/worker-rpc.h b/src/workerd/api/worker-rpc.h
@@ -395,7 +395,8 @@ class JsRpcSessionCustomEventImpl final: public WorkerInterface::CustomEvent {
 
   kj::Promise<Result> run(
       kj::Own<IoContext::IncomingRequest> incomingRequest,
-      kj::Maybe<kj::StringPtr> entrypointName) override;
+      kj::Maybe<kj::StringPtr> entrypointName,
+      kj::TaskSet& waitUntilTasks) override;
 
   kj::Promise<Result> sendRpc(
       capnp::HttpOverCapnpFactory& httpOverCapnpFactory,

diff --git a/src/workerd/io/worker-entrypoint.c++ b/src/workerd/io/worker-entrypoint.c++
@@ -650,7 +650,8 @@ kj::Promise<WorkerInterface::CustomEvent::Result>
   this->incomingRequest = kj::none;
 
   auto& context = incomingRequest->getContext();
-  auto promise = event->run(kj::mv(incomingRequest), entrypointName).attach(kj::mv(event));
+  auto promise = event->run(kj::mv(incomingRequest), entrypointName, waitUntilTasks)
+      .attach(kj::mv(event));
 
   // TODO(cleanup): In theory `context` may have been destroyed by now if `event->run()` dropped
   //   the `incomingRequest` synchronously. No current implementation does that, and

diff --git a/src/workerd/io/worker-interface.h b/src/workerd/io/worker-interface.h
@@ -106,7 +106,8 @@ class WorkerInterface: public kj::HttpService {
     // for this event.
     virtual kj::Promise<Result> run(
         kj::Own<IoContext_IncomingRequest> incomingRequest,
-        kj::Maybe<kj::StringPtr> entrypointName) = 0;
+        kj::Maybe<kj::StringPtr> entrypointName,
+        kj::TaskSet& waitUntilTasks) = 0;
 
     // Forward the event over RPC.
     virtual kj::Promise<Result> sendRpc(