Merge pull request #2062 from cloudflare/kenton/rpc-output-gate

JSRPC: Honor output gates.

Author: kentonv

src/workerd/api/tests/js-rpc-test.js

@@ -412,24 +412,33 @@ export let basicServiceBinding = {
     assert.strictEqual(await env.self.oneArg(3), 36);
     assert.strictEqual(await env.self.oneArgOmitCtx(3), 37);
     assert.strictEqual(await env.self.oneArgOmitEnvCtx(3), 6);
-    assert.rejects(() => env.self.twoArgs(123, 2),
+    await assert.rejects(() => env.self.twoArgs(123, 2), {
+      name: "TypeError",
+      message:
         "Cannot call handler function \"twoArgs\" over RPC because it has the wrong " +
         "number of arguments. A simple function handler can only be called over RPC if it has " +
         "exactly the arguments (arg, env, ctx), where only the first argument comes from the " +
         "client. To support multi-argument RPC functions, use class-based syntax (extending " +
-        "WorkerEntrypoint) instead.");
-    assert.rejects(() => env.self.noArgs(),
+        "WorkerEntrypoint) instead."
+    });
+    await assert.rejects(() => env.self.noArgs(), {
+      name: "TypeError",
+      message:
         "Attempted to call RPC function \"noArgs\" with the wrong number of arguments. " +
         "When calling a top-level handler function that is not declared as part of a class, you " +
         "must always send exactly one argument. In order to support variable numbers of " +
         "arguments, the server must use class-based syntax (extending WorkerEntrypoint) " +
-        "instead.");
-    assert.rejects(() => env.self.oneArg(1, 2),
+        "instead."
+    });
+    await assert.rejects(() => env.self.oneArg(1, 2), {
+      name: "TypeError",
+      message:
         "Attempted to call RPC function \"oneArg\" with the wrong number of arguments. " +
         "When calling a top-level handler function that is not declared as part of a class, you " +
         "must always send exactly one argument. In order to support variable numbers of " +
         "arguments, the server must use class-based syntax (extending WorkerEntrypoint) " +
-        "instead.");
+        "instead."
+    });
 
     // If we restore multi-arg support, remove the `rejects` checks above and un-comment these:
     // assert.strictEqual(await env.self.noArgs(), 13);

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

@@ -467,6 +467,12 @@ JsRpcPromiseAndPipleine callImpl(
 
       auto& ioContext = IoContext::current();
 
+      KJ_IF_SOME(lock, ioContext.waitForOutputLocksIfNecessary()) {
+        // Replace the client with a promise client that will delay thecall until the output gate
+        // is open.
+        client = lock.then([client = kj::mv(client)]() mutable { return kj::mv(client); });
+      }
+
       auto builder = client.callRequest();
 
       // This code here is slightly overcomplicated in order to avoid pushing anything to the
@@ -939,9 +945,9 @@ public:
     // object's lifetime is that of the RPC call, but in reality they are refcounted under the
     // hood. Since well be executing the call in the JS microtask queue, we have no ability to
     // actually cancel execution if a cancellation arrives over RPC, and at the end of that
-    // execution we're going to accell the call context to write the results. We could invent some
+    // execution we're going to access the call context to write the results. We could invent some
     // complicated way to skip initializing results in the case the call has been canceled, but
-    // it's easier and safer to just grap a refcount on the call context object itself, which
+    // it's easier and safer to just grab a refcount on the call context object itself, which
     // fully protects us. So... do that.
     auto ownCallContext = capnp::CallContextHook::from(callContext).addRef();
 
@@ -1052,7 +1058,13 @@ public:
           js.throwException(kj::mv(error));
         })));
 
-        return result;
+        if (ctx.hasOutputGate()) {
+          return result.then([this]() {
+            return KJ_REQUIRE_NONNULL(weakIoContext->tryGet()).waitForOutputLocks();
+          });
+        } else {
+          return result;
+        }
       };
 
       switch (op.which()) {