Implement AbortSignal serialization via ExternalPusher.

Author: kentonv

src/workerd/api/basics.c++

@@ -574,6 +574,10 @@ class AbortTriggerRpcClient final {
 
 namespace {
 // The jsrpc handler that receives aborts from the remote and triggers them locally
+//
+// TODO(cleanup): This class has been copied to external-pusher.c++. The copy here can be
+//   deleted as soon as we've switched from StreamSink to ExternalPusher and can delete all the
+//   StreamSink-related code. For now I'm not trying to avoid duplication.
 class AbortTriggerRpcServer final: public rpc::AbortTrigger::Server {
  public:
   AbortTriggerRpcServer(kj::Own<kj::PromiseFulfiller<void>> fulfiller,
@@ -858,15 +862,28 @@ void AbortSignal::serialize(jsg::Lock& js, jsg::Serializer& serializer) {
     return;
   }
 
-  auto streamCap = externalHandler
-                       ->writeStream([&](rpc::JsValue::External::Builder builder) mutable {
-    builder.setAbortTrigger();
-  }).castAs<rpc::AbortTrigger>();
+  auto triggerCap = [&]() -> rpc::AbortTrigger::Client {
+    KJ_IF_SOME(pusher, externalHandler->getExternalPusher()) {
+      auto pipeline = pusher.pushAbortSignalRequest(capnp::MessageSize{2, 0}).sendForPipeline();
+
+      externalHandler->write(
+          [signal = pipeline.getSignal()](rpc::JsValue::External::Builder builder) mutable {
+        builder.setAbortSignal(kj::mv(signal));
+      });
+
+      return pipeline.getTrigger();
+    } else {
+      return externalHandler
+          ->writeStream([&](rpc::JsValue::External::Builder builder) mutable {
+        builder.setAbortTrigger();
+      }).castAs<rpc::AbortTrigger>();
+    }
+  }();
 
   auto& ioContext = IoContext::current();
   // Keep track of every AbortSignal cloned from this one.
   // If this->triggerAbort(...) is called, each rpcClient will be informed.
-  rpcClients.add(ioContext.addObject(kj::heap<AbortTriggerRpcClient>(kj::mv(streamCap))));
+  rpcClients.add(ioContext.addObject(kj::heap<AbortTriggerRpcClient>(kj::mv(triggerCap))));
 }
 
 jsg::Ref<AbortSignal> AbortSignal::deserialize(
@@ -890,20 +907,41 @@ jsg::Ref<AbortSignal> AbortSignal::deserialize(
     return js.alloc<AbortSignal>(/* exception */ kj::none, /* maybeReason */ kj::none, flag);
   }
 
-  auto reader = externalHandler->read();
-  KJ_REQUIRE(reader.isAbortTrigger(), "external table slot type does't match serialization tag");
-
   // The AbortSignalImpl will receive any remote triggerAbort requests and fulfill the promise with the reason for abort
 
   auto signal = js.alloc<AbortSignal>(/* exception */ kj::none, /* maybeReason */ kj::none, flag);
 
-  auto paf = kj::newPromiseAndFulfiller<void>();
-  auto pendingReason = IoContext::current().addObject(kj::refcounted<PendingReason>());
+  auto& ioctx = IoContext::current();
 
-  externalHandler->setLastStream(
-      kj::heap<AbortTriggerRpcServer>(kj::mv(paf.fulfiller), kj::addRef(*pendingReason)));
-  signal->rpcAbortPromise = IoContext::current().addObject(kj::heap(kj::mv(paf.promise)));
-  signal->pendingReason = kj::mv(pendingReason);
+  auto reader = externalHandler->read();
+  if (reader.isAbortTrigger()) {
+    // Old-style StreamSink.
+    // TODO(cleanup): Remove this once the ExternalPusher autogate has rolled out.
+    auto paf = kj::newPromiseAndFulfiller<void>();
+    auto pendingReason = ioctx.addObject(kj::refcounted<PendingReason>());
+
+    externalHandler->setLastStream(
+        kj::heap<AbortTriggerRpcServer>(kj::mv(paf.fulfiller), kj::addRef(*pendingReason)));
+    signal->rpcAbortPromise = ioctx.addObject(kj::heap(kj::mv(paf.promise)));
+    signal->pendingReason = kj::mv(pendingReason);
+  } else {
+    KJ_REQUIRE(reader.isAbortSignal(), "external table slot type does't match serialization tag");
+
+    // TODO(now): This is very ugly, but it'll get better in a later commit in this PR when we
+    //   make ExternalPusher resolve all promises before delivery.
+    auto promise =
+        ioctx.getExternalPusher()
+            ->unwrapAbortSignal(reader.getAbortSignal())
+            .then([&ioctx, &signal = *signal](ExternalPusherImpl::AbortSignal resolvedSignal) {
+      return ioctx.run(
+          [&ioctx, &signal, resolvedSignal = kj::mv(resolvedSignal)](auto& lock) mutable {
+        signal.pendingReason = ioctx.addObject(kj::mv(resolvedSignal.reason));
+        return kj::mv(resolvedSignal.signal);
+      });
+    });
+
+    signal->rpcAbortPromise = ioctx.addObject(kj::heap(kj::mv(promise)));
+  }
 
   return signal;
 }

src/workerd/api/basics.h

@@ -8,6 +8,7 @@
 // TODO(cleanup): Rename to events.h?
 
 #include <workerd/io/compatibility-date.capnp.h>
+#include <workerd/io/external-pusher.h>
 #include <workerd/io/io-own.h>
 #include <workerd/io/worker-interface.capnp.h>
 #include <workerd/jsg/jsg.h>
@@ -571,9 +572,7 @@ class AbortSignal final: public EventTarget {
       jsg::Optional<jsg::JsRef<jsg::JsValue>> maybeReason = kj::none,
       Flag flag = Flag::NONE);
 
-  using PendingReason = kj::RefcountedWrapper<
-      kj::OneOf<kj::Array<kj::byte> /* v8Serialized */, kj::Exception /* if capability is dropped */
-          >>;
+  using PendingReason = ExternalPusherImpl::PendingAbortReason;
 
   // The AbortSignal explicitly does not expose a constructor(). It is
   // illegal for user code to create an AbortSignal directly.

src/workerd/io/external-pusher.c++

@@ -139,4 +139,89 @@ kj::Promise<kj::Own<kj::AsyncInputStream>> ExternalPusherImpl::unwrapStream(
       "pushed byte stream has already been consumed");
 }
 
+// =======================================================================================
+// AbortSignal handling
+
+namespace {
+
+// The jsrpc handler that receives aborts from the remote and triggers them locally
+//
+// TODO(cleanup): This class has been copied from external-pusher.c++. The copy there can be
+//   deleted as soon as we've switched from StreamSink to ExternalPusher and can delete all the
+//   StreamSink-related code. For now I'm not trying to avoid duplication.
+class AbortTriggerRpcServer final: public rpc::AbortTrigger::Server {
+ public:
+  AbortTriggerRpcServer(kj::Own<kj::PromiseFulfiller<void>> fulfiller,
+      kj::Own<ExternalPusherImpl::PendingAbortReason>&& pendingReason)
+      : fulfiller(kj::mv(fulfiller)),
+        pendingReason(kj::mv(pendingReason)) {}
+
+  kj::Promise<void> abort(AbortContext abortCtx) override {
+    auto params = abortCtx.getParams();
+    auto reason = params.getReason().getV8Serialized();
+
+    pendingReason->getWrapped() = kj::heapArray(reason.asBytes());
+    fulfiller->fulfill();
+    return kj::READY_NOW;
+  }
+
+  kj::Promise<void> release(ReleaseContext releaseCtx) override {
+    released = true;
+    return kj::READY_NOW;
+  }
+
+  ~AbortTriggerRpcServer() noexcept(false) {
+    if (pendingReason->getWrapped() != nullptr) {
+      // Already triggered
+      return;
+    }
+
+    if (!released) {
+      pendingReason->getWrapped() = JSG_KJ_EXCEPTION(FAILED, DOMAbortError,
+          "An AbortSignal received over RPC was implicitly aborted because the connection back to "
+          "its trigger was lost.");
+    }
+
+    // Always fulfill the promise in case the AbortSignal was waiting
+    fulfiller->fulfill();
+  }
+
+ private:
+  kj::Own<kj::PromiseFulfiller<void>> fulfiller;
+  kj::Own<ExternalPusherImpl::PendingAbortReason> pendingReason;
+  bool released = false;
+};
+
+}  // namespace
+
+class ExternalPusherImpl::AbortSignalImpl final: public ExternalPusher::AbortSignal::Server {
+ public:
+  AbortSignalImpl(AbortSignal content): content(kj::mv(content)) {}
+
+  kj::Maybe<AbortSignal> content;
+};
+
+kj::Promise<void> ExternalPusherImpl::pushAbortSignal(PushAbortSignalContext context) {
+  auto paf = kj::newPromiseAndFulfiller<void>();
+  auto pendingReason = kj::refcounted<PendingAbortReason>();
+
+  auto results = context.initResults(capnp::MessageSize{4, 2});
+
+  results.setTrigger(
+      kj::heap<AbortTriggerRpcServer>(kj::mv(paf.fulfiller), kj::addRef(*pendingReason)));
+  results.setSignal(abortSignalSet.add(
+      kj::heap<AbortSignalImpl>(AbortSignal{kj::mv(paf.promise), kj::mv(pendingReason)})));
+
+  return kj::READY_NOW;
+}
+
+kj::Promise<ExternalPusherImpl::AbortSignal> ExternalPusherImpl::unwrapAbortSignal(
+    ExternalPusher::AbortSignal::Client cap) {
+  auto& unwrapped = KJ_REQUIRE_NONNULL(
+      co_await abortSignalSet.getLocalServer(cap), "pushed external is not an AbortSignal");
+
+  co_return KJ_REQUIRE_NONNULL(kj::mv(kj::downcast<AbortSignalImpl>(unwrapped).content),
+      "pushed AbortSignal has already been consumed");
+}
+
 }  // namespace workerd

src/workerd/io/external-pusher.h

@@ -26,15 +26,33 @@ class ExternalPusherImpl: public rpc::JsValue::ExternalPusher::Server, public kj
 
   kj::Promise<kj::Own<kj::AsyncInputStream>> unwrapStream(ExternalPusher::InputStream::Client cap);
 
+  // Box which holds the reason why an AbortSignal was aborted. May be either:
+  // - A serialized V8 value if the signal was aborted from JavaScript.
+  // - A KJ exception if the connection from the trigger was lost.
+  using PendingAbortReason = kj::RefcountedWrapper<kj::OneOf<kj::Array<kj::byte>, kj::Exception>>;
+
+  struct AbortSignal {
+    // Resolves when `reason` has been filled in.
+    kj::Promise<void> signal;
+
+    // The abort reason box, will be uninitialized until `signal` resolves.
+    kj::Own<PendingAbortReason> reason;
+  };
+
+  kj::Promise<AbortSignal> unwrapAbortSignal(ExternalPusher::AbortSignal::Client cap);
+
  protected:
   kj::Promise<void> pushByteStream(PushByteStreamContext context) override;
+  kj::Promise<void> pushAbortSignal(PushAbortSignalContext context) override;
 
  private:
   capnp::ByteStreamFactory& byteStreamFactory;
 
   capnp::CapabilityServerSet<ExternalPusher::InputStream> inputStreamSet;
+  capnp::CapabilityServerSet<ExternalPusher::AbortSignal> abortSignalSet;
 
   class InputStreamImpl;
+  class AbortSignalImpl;
 };
 
 }  // namespace workerd