Merge pull request #5862 from cloudflare/felix/010926-stw-cancel-revert

Revert "Remove block on tail stream capability drop."

Author: fhanau

src/workerd/io/trace-stream.c++

@@ -976,21 +976,24 @@ kj::Promise<WorkerInterface::CustomEvent::Result> TailStreamCustomEvent::sendRpc
 
   capFulfiller->fulfill(kj::mv(cap));
 
+  // Forked promise for completion of all capabilities associated with the cap stream. This is
+  // expected to be resolved when the request is canceled or when the client receives the stop
+  // signal and deallocates cap after the tail worker indicates that it has processed all events
+  // successfully.
+  kj::ForkedPromise<void> forked = completionPaf.promise.fork();
   try {
-    // Wait for EITHER:
-    // 1. The RPC to complete successfully (returns the event outcome), OR
-    // 2. The capability to be dropped by the client (returns CANCELED)
-    // Whichever happens first determines our result.
     EventOutcome outcome = co_await sent.then([](auto resp) {
       return resp.getResult();
-    }).exclusiveJoin(completionPaf.promise.then([]() { return EventOutcome::CANCELED; }));
+    }).exclusiveJoin(forked.addBranch().then([]() { return EventOutcome::CANCELED; }));
 
+    // If the sent promise returned first, we still need to wait for the parent process to drop the
+    // capability (which should happen right after it receives the stop signal) so that no
+    // capabilities remain in an incomplete state when we return.
+    co_await forked.addBranch();
     co_return WorkerInterface::CustomEvent::Result{.outcome = outcome};
   } catch (...) {
-    // If an exception occurs, capture it and ensure proper cleanup
     auto e = kj::getCaughtExceptionAsKj();
     if (revokePaf.fulfiller->isWaiting()) {
-      // Reject the revoke promise to trigger capability revocation
       revokePaf.fulfiller->reject(kj::cp(e));
     }
     kj::throwFatalException(kj::mv(e));