Wait for pushed externals upfront.

This commit refactors ExternalPusher a bit such that we wait for all the pushed externals in an external set upfront, before deserializing the value.

This has two benefits:
* It removes some complexity from both the ReadableStream and AbortSignal deserializers.
* It sets us up to support embedded promise pipelining later -- that is, embedding promises into another message, and having the promises awaited and replaced with their resolution before delivery. This is already supported by Cap'n Web, and we want to support it in workerd's built-in RPC as well to match.

Author: kentonv

src/workerd/api/basics.c++

@@ -927,20 +927,11 @@ jsg::Ref<AbortSignal> AbortSignal::deserialize(
   } 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);
-      });
-    });
+    auto resolvedSignal =
+        externalHandler->getPushedExternals().unwrapAbortSignal(reader.getAbortSignal());
 
-    signal->rpcAbortPromise = ioctx.addObject(kj::heap(kj::mv(promise)));
+    signal->rpcAbortPromise = ioctx.addObject(kj::heap(kj::mv(resolvedSignal.signal)));
+    signal->pendingReason = ioctx.addObject(kj::mv(resolvedSignal.reason));
   }
 
   return signal;

src/workerd/api/streams/readable.c++

@@ -743,9 +743,7 @@ jsg::Ref<ReadableStream> ReadableStream::deserialize(
 
   kj::Own<kj::AsyncInputStream> in;
   if (rs.hasStream()) {
-    // TODO(now): Create a mechanism to await all the external promises *before* deserializing.
-    //   We'll need it anyway to support Promise serialization.
-    in = kj::newPromisedStream(ioctx.getExternalPusher()->unwrapStream(rs.getStream()));
+    in = externalHandler->getPushedExternals().unwrapStream(rs.getStream());
   } else {
     kj::Maybe<uint64_t> expectedLength;
     auto el = rs.getExpectedLength();

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

@@ -217,11 +217,14 @@ struct DeserializeResult {
 };
 
 // Call to construct a JS value from an `rpc::JsValue`.
-DeserializeResult deserializeJsValue(
-    jsg::Lock& js, rpc::JsValue::Reader reader, kj::Maybe<StreamSinkImpl&> streamSink = kj::none) {
+DeserializeResult deserializeJsValue(jsg::Lock& js,
+    rpc::JsValue::Reader reader,
+    kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals,
+    kj::Maybe<StreamSinkImpl&> streamSink = kj::none) {
   auto disposalGroup = kj::heap<RpcStubDisposalGroup>();
 
-  RpcDeserializerExternalHandler externalHandler(reader.getExternals(), *disposalGroup, streamSink);
+  RpcDeserializerExternalHandler externalHandler(
+      reader.getExternals(), kj::mv(pushedExternals), *disposalGroup, streamSink);
 
   jsg::Deserializer deserializer(js, reader.getV8Serialized(), kj::none, kj::none,
       jsg::Deserializer::Options{
@@ -249,8 +252,10 @@ DeserializeResult deserializeJsValue(
 // an object) which disposes all stubs therein.
 jsg::JsValue deserializeRpcReturnValue(jsg::Lock& js,
     rpc::JsRpcTarget::CallResults::Reader callResults,
+    kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals,
     kj::Maybe<StreamSinkImpl&> streamSink) {
-  auto [value, disposalGroup, ss] = deserializeJsValue(js, callResults.getResult(), streamSink);
+  auto [value, disposalGroup, ss] =
+      deserializeJsValue(js, callResults.getResult(), kj::mv(pushedExternals), streamSink);
 
   if (streamSink == kj::none) {
     KJ_REQUIRE(ss == kj::none,
@@ -549,12 +554,13 @@ JsRpcPromiseAndPipeline callImpl(jsg::Lock& js,
       }
 
       kj::Maybe<kj::Own<StreamSinkImpl>> resultStreamSink;
+      auto externalPusher = ioContext.getExternalPusher();
       if (useExternalPusher) {
         // Unfortunately, we always have to send the ExternalPusher since we don't know whether the
         // call will return any streams (or other pushed externals). Luckily, it's a
         // one-per-IoContext object, not a big deal. (It'll take a slot on the capnp export table
         // though.)
-        builder.getResultsStreamHandler().setExternalPusher(ioContext.getExternalPusher());
+        builder.getResultsStreamHandler().setExternalPusher(externalPusher.addRef());
       } else {
         // Unfortunately, we always have to send a `resultsStreamSink` because we don't know until
         // after the call completes whether or not it will return any streams. If it's unused,
@@ -569,6 +575,31 @@ JsRpcPromiseAndPipeline callImpl(jsg::Lock& js,
         ssf->fulfill(callResult.getParamsStreamSink());
       }
 
+      struct ResponseAndPushedExternals {
+        capnp::Response<rpc::JsRpcTarget::CallResults> response;
+        kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals;
+      };
+
+      auto responsePromise =
+          callResult.then([externalPusher = kj::mv(externalPusher)](
+                              capnp::Response<rpc::JsRpcTarget::CallResults> response) mutable
+              -> kj::Promise<ResponseAndPushedExternals> {
+        KJ_IF_SOME(promise, externalPusher->makeBundle(response.getResult().getExternals())) {
+          return promise.then([response = kj::mv(response)](
+                                  kj::Own<ExternalPusherImpl::Bundle> pushedExternals) mutable {
+            return ResponseAndPushedExternals{
+              .response = kj::mv(response),
+              .pushedExternals = kj::mv(pushedExternals),
+            };
+          });
+        } else {
+          return ResponseAndPushedExternals{
+            .response = kj::mv(response),
+            .pushedExternals = kj::none,
+          };
+        }
+      });
+
       // We need to arrange that our JsRpcPromise will updated in-place with the final settlement
       // of this RPC promise. However, we can't actually construct the JsRpcPromise until we have
       // the final promise to give it. To resolve the cycle, we only create a JsRpcPromise::WeakRef
@@ -577,11 +608,11 @@ JsRpcPromiseAndPipeline callImpl(jsg::Lock& js,
 
       // RemotePromise lets us consume its pipeline and promise portions independently; we consume
       // the promise here and we consume the pipeline below, both via kj::mv().
-      auto jsPromise = ioContext.awaitIo(js, kj::mv(callResult),
+      auto jsPromise = ioContext.awaitIo(js, kj::mv(responsePromise),
           [weakRef = kj::atomicAddRef(*weakRef), resultStreamSink = kj::mv(resultStreamSink)](
-              jsg::Lock& js,
-              capnp::Response<rpc::JsRpcTarget::CallResults> response) mutable -> jsg::Value {
-        auto jsResult = deserializeRpcReturnValue(js, response, resultStreamSink);
+              jsg::Lock& js, ResponseAndPushedExternals response) mutable -> jsg::Value {
+        auto jsResult = deserializeRpcReturnValue(
+            js, kj::mv(response.response), kj::mv(response.pushedExternals), resultStreamSink);
 
         if (weakRef->disposed) {
           // The promise was explicitly disposed before it even resolved. This means we must dispose
@@ -973,19 +1004,23 @@ class JsRpcTargetBase: public rpc::JsRpcTarget::Server {
   // or during a call.
   JsRpcTargetBase(IoContext& ctx, MayOutliveIncomingRequest)
       : enterIsolateAndCall(ctx.makeReentryCallback<IoContext::TOP_UP>(
-            [this, &ctx](Worker::Lock& lock, CallContext callContext) {
-              return callImpl(lock, ctx, callContext);
+            [this, &ctx](Worker::Lock& lock,
+                CallContext callContext,
+                kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals) {
+              return callImpl(lock, ctx, callContext, kj::mv(pushedExternals));
             })),
         externalPusher(ctx.getExternalPusher()) {}
 
   // Constructor use by EntrypointJsRpcTarget, which is revoked and destroyed before the IoContext
   // can possibly be canceled. It can just use ctx.run().
   JsRpcTargetBase(IoContext& ctx, CantOutliveIncomingRequest)
-      : enterIsolateAndCall([this, &ctx](CallContext callContext) {
+      : enterIsolateAndCall([this, &ctx](CallContext callContext,
+                                kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals) {
           // Note: No need to topUpActor() since this is the start of a top-level request, so the
           // actor will already have been topped up by IncomingRequest::delivered().
-          return ctx.run([this, &ctx, callContext](Worker::Lock& lock) mutable {
-            return callImpl(lock, ctx, callContext);
+          return ctx.run([this, &ctx, callContext, pushedExternals = kj::mv(pushedExternals)](
+                             Worker::Lock& lock) mutable {
+            return callImpl(lock, ctx, callContext, kj::mv(pushedExternals));
           });
         }),
         externalPusher(ctx.getExternalPusher()) {}
@@ -1012,8 +1047,25 @@ class JsRpcTargetBase: public rpc::JsRpcTarget::Server {
 
   // Handles the delivery of JS RPC method calls.
   kj::Promise<void> call(CallContext callContext) override {
-    // Try to execute the requested method.
-    return enterIsolateAndCall(callContext).catch_([](kj::Exception&& e) {
+    kj::Promise<void> promise = [&]() {
+      auto op = callContext.getParams().getOperation();
+      if (op.hasCallWithArgs()) {
+        auto args = op.getCallWithArgs();
+        if (args.hasExternals()) {
+          KJ_IF_SOME(bundlePromise, externalPusher->makeBundle(args.getExternals())) {
+            // The call params contain one or more pushed externals.
+            return bundlePromise.then(
+                [this, callContext](kj::Own<ExternalPusherImpl::Bundle> pushedExternals) mutable {
+              return enterIsolateAndCall(callContext, kj::mv(pushedExternals));
+            });
+          }
+        }
+      }
+
+      return enterIsolateAndCall(callContext, kj::none);
+    }();
+
+    return promise.catch_([](kj::Exception&& e) {
       if (jsg::isTunneledException(e.getDescription())) {
         // Annotate exceptions in RPC worker calls as remote exceptions.
         auto description = jsg::stripRemoteExceptionPrefix(e.getDescription());
@@ -1039,14 +1091,18 @@ class JsRpcTargetBase: public rpc::JsRpcTarget::Server {
 
   // Function which enters the isolate lock and IoContext and then invokes callImpl(). Created
   // using IoContext::makeReentryCallback().
-  kj::Function<kj::Promise<void>(CallContext callContext)> enterIsolateAndCall;
+  kj::Function<kj::Promise<void>(CallContext, kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>>)>
+      enterIsolateAndCall;
 
   kj::Rc<ExternalPusherImpl> externalPusher;
 
   // Returns true if the given name cannot be used as a method on this type.
   virtual bool isReservedName(kj::StringPtr name) = 0;
 
-  kj::Promise<void> callImpl(Worker::Lock& lock, IoContext& ctx, CallContext callContext) {
+  kj::Promise<void> callImpl(Worker::Lock& lock,
+      IoContext& ctx,
+      CallContext callContext,
+      kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals) {
     jsg::Lock& js = lock;
     auto params = callContext.getParams();
     // Method name suitable for use in trace and error messages. May be a pointer into the RPC
@@ -1236,10 +1292,10 @@ class JsRpcTargetBase: public rpc::JsRpcTarget::Server {
 
         InvocationResult invocationResult;
         KJ_IF_SOME(envCtx, targetInfo.envCtx) {
-          invocationResult = invokeFnInsertingEnvCtx(
-              js, methodNameForTrace, fn, thisArg, args, envCtx.env, envCtx.ctx);
+          invocationResult = invokeFnInsertingEnvCtx(js, methodNameForTrace, fn, thisArg, args,
+              kj::mv(pushedExternals), envCtx.env, envCtx.ctx);
         } else {
-          invocationResult = invokeFn(js, fn, thisArg, args);
+          invocationResult = invokeFn(js, fn, thisArg, args, kj::mv(pushedExternals));
         }
 
         // We have a function, so let's call it and serialize the result for RPC.
@@ -1406,10 +1462,11 @@ class JsRpcTargetBase: public rpc::JsRpcTarget::Server {
   static InvocationResult invokeFn(jsg::Lock& js,
       v8::Local<v8::Function> fn,
       v8::Local<v8::Object> thisArg,
-      kj::Maybe<rpc::JsValue::Reader> args) {
+      kj::Maybe<rpc::JsValue::Reader> args,
+      kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals) {
     // We received arguments from the client, deserialize them back to JS.
     KJ_IF_SOME(a, args) {
-      auto [value, disposalGroup, streamSink] = deserializeJsValue(js, a);
+      auto [value, disposalGroup, streamSink] = deserializeJsValue(js, a, kj::mv(pushedExternals));
       auto args = KJ_REQUIRE_NONNULL(
           value.tryCast<jsg::JsArray>(), "expected JsArray when deserializing arguments.");
       // Call() expects a `Local<Value> []`... so we populate an array.
@@ -1440,6 +1497,7 @@ class JsRpcTargetBase: public rpc::JsRpcTarget::Server {
       v8::Local<v8::Function> fn,
       v8::Local<v8::Object> thisArg,
       kj::Maybe<rpc::JsValue::Reader> args,
+      kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals,
       v8::Local<v8::Value> env,
       jsg::JsObject ctx) {
     // Determine the function arity (how many parameters it was declared to accept) by reading the
@@ -1469,7 +1527,7 @@ class JsRpcTargetBase: public rpc::JsRpcTarget::Server {
     kj::Maybe<jsg::JsArray> argsArrayFromClient;
     size_t argCountFromClient = 0;
     KJ_IF_SOME(a, args) {
-      auto [value, disposalGroup, ss] = deserializeJsValue(js, a);
+      auto [value, disposalGroup, ss] = deserializeJsValue(js, a, kj::mv(pushedExternals));
       streamSink = kj::mv(ss);
 
       auto array = KJ_REQUIRE_NONNULL(

src/workerd/api/worker-rpc.h

@@ -127,16 +127,25 @@ class RpcDeserializerExternalHandler final: public jsg::Deserializer::ExternalHa
   // The `streamSink` parameter should be provided if a StreamSink already exists, e.g. when
   // deserializing results. If omitted, it will be constructed on-demand.
   RpcDeserializerExternalHandler(capnp::List<rpc::JsValue::External>::Reader externals,
+      kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals,
       RpcStubDisposalGroup& disposalGroup,
       kj::Maybe<StreamSinkImpl&> streamSink)
       : externals(externals),
+        pushedExternals(kj::mv(pushedExternals)),
         disposalGroup(disposalGroup),
         streamSink(streamSink) {}
   ~RpcDeserializerExternalHandler() noexcept(false);
 
   // Read and return the next external.
   rpc::JsValue::External::Reader read();
 
+  // Call to get the bundle of externals pushed to this isolate using ExternalPusher. Throws an
+  // exception if there were none -- but this should only be called when a specific pushed
+  // external was encountered during deserialization.
+  ExternalPusherImpl::Bundle& getPushedExternals() {
+    return *KJ_REQUIRE_NONNULL(pushedExternals);
+  }
+
   // Call immediately after `read()` when reading an external that is associated with a stream.
   // `stream` is published back to the sender via StreamSink.
   void setLastStream(capnp::Capability::Client stream);
@@ -158,6 +167,8 @@ class RpcDeserializerExternalHandler final: public jsg::Deserializer::ExternalHa
   capnp::List<rpc::JsValue::External>::Reader externals;
   uint i = 0;
 
+  kj::Maybe<kj::Own<ExternalPusherImpl::Bundle>> pushedExternals;
+
   kj::UnwindDetector unwindDetector;
   RpcStubDisposalGroup& disposalGroup;
 

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

@@ -7,6 +7,70 @@
 
 namespace workerd {
 
+// HACK: If `CapabilityServerSet` successfully unwraps a `ClientHook`, this either means the hook
+//   itself pointed to a local object, or it resolved to something that did. We want to populate
+//   a HashMap based on hook pointers, so we'd better use the fully-resolved hook, otherwise things
+//   like membrane layers can cause confusion.
+static capnp::ClientHook* fullyResolvedPtr(capnp::ClientHook& hook) {
+  KJ_IF_SOME(r, hook.getResolved()) {
+    return &r;
+  } else {
+    return &hook;
+  }
+}
+
+kj::Maybe<kj::Promise<kj::Own<ExternalPusherImpl::Bundle>>> ExternalPusherImpl::makeBundle(
+    capnp::List<rpc::JsValue::External>::Reader externals) {
+  kj::Vector<kj::Promise<void>> promises;
+  kj::Maybe<kj::Own<Bundle>> maybeBundle;
+
+  auto add = [&](auto cap, auto func) {
+    auto hook = capnp::ClientHook::from(cap);
+    Bundle* bundle;
+    KJ_IF_SOME(b, maybeBundle) {
+      bundle = b.get();
+    } else {
+      bundle = maybeBundle.emplace(kj::heap<Bundle>());
+    }
+    promises.add((this->*func)(cap).then([bundle, hook = kj::mv(hook)](auto&& unwrapped) mutable {
+      bundle->items.insert(fullyResolvedPtr(*hook), kj::mv(unwrapped));
+    }));
+  };
+
+  for (auto ext: externals) {
+    switch (ext.which()) {
+      case rpc::JsValue::External::READABLE_STREAM: {
+        auto rs = ext.getReadableStream();
+        if (rs.hasStream()) {
+          add(rs.getStream(), &ExternalPusherImpl::unwrapStream);
+        }
+        break;
+      }
+
+      case rpc::JsValue::External::ABORT_SIGNAL:
+        add(ext.getAbortSignal(), &ExternalPusherImpl::unwrapAbortSignal);
+        break;
+
+      default:
+        break;
+    }
+  }
+
+  return maybeBundle.map([&](kj::Own<Bundle>& bundle) -> kj::Promise<kj::Own<Bundle>> {
+    return kj::joinPromisesFailFast(promises.releaseAsArray())
+        .then([bundle = kj::mv(bundle)]() mutable { return kj::mv(bundle); });
+  });
+}
+
+template <typename T>
+T ExternalPusherImpl::Bundle::unwrap(capnp::Capability::Client cap) {
+  auto hook = capnp::ClientHook::from(kj::mv(cap));
+  auto& item = KJ_REQUIRE_NONNULL(items.find(fullyResolvedPtr(*hook)));
+  auto result = kj::mv(KJ_REQUIRE_NONNULL(item.tryGet<T>()));
+  item = {};
+  return result;
+}
+
 // =======================================================================================
 // ReadableStream handling
 
@@ -139,6 +203,11 @@ kj::Promise<kj::Own<kj::AsyncInputStream>> ExternalPusherImpl::unwrapStream(
       "pushed byte stream has already been consumed");
 }
 
+kj::Own<kj::AsyncInputStream> ExternalPusherImpl::Bundle::unwrapStream(
+    ExternalPusher::InputStream::Client cap) {
+  return unwrap<kj::Own<kj::AsyncInputStream>>(kj::mv(cap));
+}
+
 // =======================================================================================
 // AbortSignal handling
 
@@ -224,4 +293,9 @@ kj::Promise<ExternalPusherImpl::AbortSignal> ExternalPusherImpl::unwrapAbortSign
       "pushed AbortSignal has already been consumed");
 }
 
+ExternalPusherImpl::AbortSignal ExternalPusherImpl::Bundle::unwrapAbortSignal(
+    ExternalPusher::AbortSignal::Client cap) {
+  return unwrap<AbortSignal>(kj::mv(cap));
+}
+
 }  // namespace workerd