Merge pull request #6233 from cloudflare/ketan/exmem-rpc

Add external memory accounting to JsRpcStub deserialization

Author: ketanhwr

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

@@ -731,10 +731,12 @@ kj::Maybe<jsg::Ref<JsRpcProperty>> JsRpcPromise::getProperty(jsg::Lock& js, kj::
   return js.alloc<JsRpcProperty>(JSG_THIS, kj::mv(name));
 }
 
-JsRpcStub::JsRpcStub(
-    IoOwn<rpc::JsRpcTarget::Client> capnpClient, RpcStubDisposalGroup& disposalGroup)
+JsRpcStub::JsRpcStub(IoOwn<rpc::JsRpcTarget::Client> capnpClient,
+    RpcStubDisposalGroup& disposalGroup,
+    jsg::ExternalMemoryAdjustment externalMemoryAdjustment)
     : capnpClient(kj::mv(capnpClient)),
-      disposalGroup(disposalGroup) {
+      disposalGroup(disposalGroup),
+      externalMemoryAdjustment(kj::mv(externalMemoryAdjustment)) {
   disposalGroup.list.add(*this);
 }
 
@@ -834,6 +836,7 @@ jsg::Ref<JsRpcStub> JsRpcStub::dup(jsg::Lock& js) {
 
 void JsRpcStub::dispose() {
   capnpClient = kj::none;
+  externalMemoryAdjustment = kj::none;
   KJ_IF_SOME(d, disposalGroup) {
     d.list.remove(*this);
     disposalGroup = kj::none;
@@ -916,8 +919,13 @@ jsg::Ref<JsRpcStub> JsRpcStub::deserialize(
   KJ_REQUIRE(reader.isRpcTarget(), "external table slot type doesn't match serialization tag");
 
   auto& ioctx = IoContext::current();
-  return js.alloc<JsRpcStub>(
-      ioctx.addObject(kj::heap(reader.getRpcTarget())), externalHandler->getDisposalGroup());
+
+  // Account for membrane/promise memory in the KJ heap (~1600 bytes per stub from profiling).
+  static constexpr size_t ESTIMATED_EXTERNAL_MEMORY_PER_STUB = 1600;
+  auto externalMemory = js.getExternalMemoryAdjustment(ESTIMATED_EXTERNAL_MEMORY_PER_STUB);
+
+  return js.alloc<JsRpcStub>(ioctx.addObject(kj::heap(reader.getRpcTarget())),
+      externalHandler->getDisposalGroup(), kj::mv(externalMemory));
 }
 
 static bool isFunctionForRpc(jsg::Lock& js, v8::Local<v8::Function> func) {

src/workerd/api/worker-rpc.h

@@ -370,8 +370,13 @@ class JsRpcProperty: public JsRpcClientProvider {
 // `JsRpcStub::sendJsRpc()`.
 class JsRpcStub: public JsRpcClientProvider {
  public:
+  // Only deserialize() needs ExternalMemoryAdjustment — it extracts a new capability from an RPC
+  // response, creating MembraneHook + forked promise allocations. The other callers (dup() and
+  // constructor()) just refcount existing capabilities or create local loopbacks.
   JsRpcStub(IoOwn<rpc::JsRpcTarget::Client> capnpClient): capnpClient(kj::mv(capnpClient)) {}
-  JsRpcStub(IoOwn<rpc::JsRpcTarget::Client> capnpClient, RpcStubDisposalGroup& disposalGroup);
+  JsRpcStub(IoOwn<rpc::JsRpcTarget::Client> capnpClient,
+      RpcStubDisposalGroup& disposalGroup,
+      jsg::ExternalMemoryAdjustment externalMemoryAdjustment);
   ~JsRpcStub() noexcept(false);
 
   rpc::JsRpcTarget::Client getClient();
@@ -414,6 +419,7 @@ class JsRpcStub: public JsRpcClientProvider {
 
   kj::Maybe<RpcStubDisposalGroup&> disposalGroup;
   kj::ListLink<JsRpcStub> disposalGroupLink;
+  kj::Maybe<jsg::ExternalMemoryAdjustment> externalMemoryAdjustment;
 
   friend class RpcStubDisposalGroup;
 };