From 87cb25624c653308d7dbc82173f5a7ba14a46391 Mon Sep 17 00:00:00 2001
From: Kenton Varda <kenton@cloudflare.com>
Date: Fri, 26 Apr 2024 13:55:35 -0500
Subject: [PATCH] JSRPC: Wait for output gates on incoming and outgoing RPCs.

(Tested in internal codebase where some additional utilities are available.)
---
 src/workerd/api/worker-rpc.c++ | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/workerd/api/worker-rpc.c++ b/src/workerd/api/worker-rpc.c++
index a26fe8cb74e..77514ea1899 100644
--- a/src/workerd/api/worker-rpc.c++
+++ b/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
@@ -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()) {