fix(Worklets): RemoteFunction memory cycle

packages/react-native-worklets/Common/cpp/worklets/Compat/StableApi.cpp

@@ -2,6 +2,7 @@
 #include <worklets/Compat/Holders.h>
 #include <worklets/Compat/StableApi.h>
 #include <worklets/SharedItems/Serializable.h>
+#include <worklets/SharedItems/SerializableRemoteFunction.h>
 #include <worklets/SharedItems/Shareable.h>
 #include <worklets/SharedItems/Synchronizable.h>
 #include <worklets/Tools/JSISerializer.h>

packages/react-native-worklets/Common/cpp/worklets/NativeModules/JSIWorkletsModuleProxy.cpp

@@ -5,6 +5,7 @@
 #include <worklets/NativeModules/JSIWorkletsModuleProxy.h>
 #include <worklets/SharedItems/Serializable.h>
 #include <worklets/SharedItems/SerializableFactory.h>
+#include <worklets/SharedItems/SerializableRemoteFunction.h>
 #include <worklets/SharedItems/Shareable.h>
 #include <worklets/SharedItems/Synchronizable.h>
 #include <worklets/Tools/FeatureFlags.h>
@@ -348,11 +349,11 @@ jsi::Object JSIWorkletsModuleProxy::toOptimizedObject(jsi::Runtime &rt) const {
         }
         if (hostRuntimeId == RuntimeData::rnRuntimeId) {
           const int remoteId = static_cast<int>(at<1>(args).getNumber());
-          auto ref = makeSerializableRemoteFunction(rt, name, remoteId, jsScheduler);
+          auto ref = makeRNOriginSerializableRemoteFunction(rt, name, remoteId, jsScheduler);
           ref.asObject(rt).setProperty(rt, "__keepAlive", true);
           return ref;
         }
-        return makeSerializableRemoteFunction(rt, name, std::move(fun), hostRuntimeId);
+        return makeWorkletOriginSerializableRemoteFunction(rt, name, std::move(fun), hostRuntimeId);
       });
 
   jsi_utils::addMethod<2>(

packages/react-native-worklets/Common/cpp/worklets/SharedItems/Serializable.cpp

@@ -7,6 +7,7 @@
 #include <worklets/WorkletRuntime/WorkletRuntime.h>
 
 #include <memory>
+#include <stdexcept>
 #include <string>
 #include <utility>
 
@@ -22,18 +23,6 @@ jsi::Function getValueUnpacker(jsi::Runtime &rt) {
   return valueUnpacker.asObject(rt).asFunction(rt);
 }
 
-jsi::Function getRemoteFunctionUnpacker(jsi::Runtime &rt) {
-  auto remoteFunctionUnpacker = rt.global().getProperty(rt, "__remoteFunctionUnpacker");
-  react_native_assert(remoteFunctionUnpacker.isObject() && "remoteFunctionUnpacker not found");
-  return remoteFunctionUnpacker.asObject(rt).asFunction(rt);
-}
-
-jsi::Object getRemoteFunctionRegistry(jsi::Runtime &rt) {
-  auto registry = rt.global().getProperty(rt, "__remoteFunctionRegistry");
-  react_native_assert(registry.isObject() && "remoteFunctionRegistry not found");
-  return registry.getObject(rt);
-}
-
 } // namespace
 
 jsi::Value makeSerializableClone(
@@ -296,61 +285,6 @@ jsi::Value SerializableImport::toJSValue(jsi::Runtime &rt) {
   return metroRequire.asObject(rt).asFunction(rt).call(rt, source_).asObject(rt).getProperty(rt, imported);
 }
 
-SerializableRemoteFunction::~SerializableRemoteFunction() {
-  if (isHostedOnRNRuntime()) {
-    // TODO: consider batching
-    const auto &data = std::get<RNRuntimeData>(runtimeData_);
-    data.jsScheduler->scheduleOnJS([id = data.remoteId](jsi::Runtime &rt) {
-      const auto registry = getRemoteFunctionRegistry(rt);
-      registry.getPropertyAsFunction(rt, "delete").callWithThis(rt, registry, jsi::Value(id));
-    });
-  } else {
-    auto &workletData = std::get<WorkletRuntimeData>(runtimeData_);
-    cleanupRuntimeAware(hostRuntime_, workletData.function);
-  }
-}
-
-jsi::Value SerializableRemoteFunction::toJSValue(jsi::Runtime &rt) {
-  if (&rt == hostRuntime_) {
-    if (isHostedOnRNRuntime()) {
-      const auto &rnData = std::get<RNRuntimeData>(runtimeData_);
-      const auto registry = getRemoteFunctionRegistry(rt);
-      return registry.getPropertyAsFunction(rt, "get").callWithThis(rt, registry, jsi::Value(rnData.remoteId));
-    } else {
-      const auto &workletData = std::get<WorkletRuntimeData>(runtimeData_);
-      return jsi::Value(rt, *workletData.function);
-    }
-  } else {
-    const auto name = name_.empty() ? jsi::Value::undefined() : jsi::String::createFromUtf8(rt, name_);
-    auto holderFunction = getRemoteFunctionUnpacker(rt).call(rt, name).getObject(rt);
-    holderFunction.setNativeState(rt, std::make_shared<SerializableJSRef>(shared_from_this()));
-    return holderFunction;
-  }
-}
-
-// TODO: generalize it and merge with other scheduling methods
-void SerializableRemoteFunction::resolveOrRejectPromise(
-    const std::shared_ptr<Serializable> &resolveValue,
-    const std::shared_ptr<RuntimeManager> &runtimeManager) {
-  if (isHostedOnRNRuntime()) {
-    const auto &data = std::get<RNRuntimeData>(runtimeData_);
-    data.jsScheduler->scheduleOnJS([resolver = shared_from_this(), resolveValue](jsi::Runtime &rt) {
-      resolver->toJSValue(rt).getObject(rt).getFunction(rt).call(rt, resolveValue->toJSValue(rt));
-    });
-  } else {
-    const auto workletRuntime = runtimeManager->getRuntime(hostRuntimeId_);
-    // NOLINTNEXTLINE(readability/braces)
-    if (!workletRuntime) [[unlikely]] {
-      // Host runtime is dead, most likely we're the last owner of the Remote Function.
-      // Do nothing.
-    } else {
-      workletRuntime->schedule([resolver = shared_from_this(), resolveValue](jsi::Runtime &rt) {
-        resolver->toJSValue(rt).getObject(rt).getFunction(rt).call(rt, resolveValue->toJSValue(rt));
-      });
-    }
-  }
-}
-
 jsi::Value SerializableInitializer::toJSValue(jsi::Runtime &rt) {
   if (remoteValue_ == nullptr) {
     auto initObj = initializer_->toJSValue(rt);

packages/react-native-worklets/Common/cpp/worklets/SharedItems/Serializable.h

@@ -7,6 +7,7 @@
 #include <worklets/WorkletRuntime/RuntimeData.h>
 
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <string>
 #include <utility>
@@ -278,71 +279,6 @@ class SerializableImport : public Serializable {
   const std::string imported_;
 };
 
-/** Forward declaration */
-class RuntimeManager;
-
-class SerializableRemoteFunction : public Serializable,
-                                   public std::enable_shared_from_this<SerializableRemoteFunction> {
- private:
-  struct RNRuntimeData {
-    const int remoteId;
-    const std::shared_ptr<JSScheduler> jsScheduler;
-  };
-
-  struct WorkletRuntimeData {
-    std::unique_ptr<jsi::Value> function;
-  };
-
-  jsi::Runtime *hostRuntime_;
-  const RuntimeData::RuntimeId hostRuntimeId_;
-  std::variant<RNRuntimeData, WorkletRuntimeData> runtimeData_;
-  const std::string name_;
-
- public:
-  /** Creates RN Runtime Remote Function. */
-  SerializableRemoteFunction(
-      jsi::Runtime &rnRuntime,
-      const std::string &name,
-      const int remoteId,
-      const std::shared_ptr<JSScheduler> &jsScheduler)
-      : Serializable(ValueType::RemoteFunctionType),
-        hostRuntime_(&rnRuntime),
-        hostRuntimeId_(RuntimeData::rnRuntimeId),
-        runtimeData_(RNRuntimeData{.remoteId = remoteId, .jsScheduler = jsScheduler}),
-        name_(name) {}
-
-  /** Creates Worklet Runtime Remote Function. */
-  SerializableRemoteFunction(
-      jsi::Runtime &workletRuntime,
-      const std::string &name,
-      jsi::Function &&function,
-      RuntimeData::RuntimeId hostRuntimeId)
-      : Serializable(ValueType::RemoteFunctionType),
-        hostRuntime_(&workletRuntime),
-        hostRuntimeId_(hostRuntimeId),
-        runtimeData_(WorkletRuntimeData{.function = std::make_unique<jsi::Value>(workletRuntime, std::move(function))}),
-        name_(name) {}
-
-  ~SerializableRemoteFunction() override;
-
-  SerializableRemoteFunction(const SerializableRemoteFunction &) = delete;
-  SerializableRemoteFunction &operator=(const SerializableRemoteFunction &) = delete;
-
-  void resolveOrRejectPromise(
-      const std::shared_ptr<Serializable> &resolveValue,
-      const std::shared_ptr<RuntimeManager> &runtimeManager);
-
-  jsi::Value toJSValue(jsi::Runtime &rt) override;
-
-  [[nodiscard]] bool isHostedOnRNRuntime() const noexcept {
-    return std::holds_alternative<RNRuntimeData>(runtimeData_);
-  }
-
-  [[nodiscard]] RuntimeData::RuntimeId getHostRuntimeId() const {
-    return hostRuntimeId_;
-  }
-};
-
 class SerializableInitializer : public Serializable {
  private:
   // We don't release the initializer since the handle can get

packages/react-native-worklets/Common/cpp/worklets/SharedItems/SerializableFactory.cpp

@@ -1,5 +1,6 @@
 #include <jsi/jsi.h>
 #include <worklets/SharedItems/SerializableFactory.h>
+#include <worklets/SharedItems/SerializableRemoteFunction.h>
 
 #include <memory>
 #include <utility>
@@ -60,22 +61,22 @@ jsi::Value makeSerializableHostFunction(
   return SerializableJSRef::newNativeStateObject(rt, serializable);
 }
 
-jsi::Value makeSerializableRemoteFunction(
+jsi::Value makeRNOriginSerializableRemoteFunction(
     jsi::Runtime &rnRuntime,
     const std::string &name,
     const int remoteId,
     const std::shared_ptr<JSScheduler> &jsScheduler) {
-  auto serializable = std::make_shared<SerializableRemoteFunction>(rnRuntime, name, remoteId, jsScheduler);
+  auto serializable = std::make_shared<SerializableRemoteFunction::RNOrigin>(rnRuntime, name, remoteId, jsScheduler);
   return SerializableJSRef::newNativeStateObject(rnRuntime, serializable);
 }
 
-jsi::Value makeSerializableRemoteFunction(
+jsi::Value makeWorkletOriginSerializableRemoteFunction(
     jsi::Runtime &workletRuntime,
     const std::string &name,
     jsi::Function &&function,
     RuntimeData::RuntimeId hostRuntimeId) {
-  auto serializable =
-      std::make_shared<SerializableRemoteFunction>(workletRuntime, name, std::move(function), hostRuntimeId);
+  auto serializable = std::make_shared<SerializableRemoteFunction::WorkletOrigin>(
+      workletRuntime, name, std::move(function), hostRuntimeId);
   return SerializableJSRef::newNativeStateObject(workletRuntime, serializable);
 }
 

packages/react-native-worklets/Common/cpp/worklets/SharedItems/SerializableFactory.h

@@ -59,15 +59,13 @@ jsi::Value makeSerializableHostFunction(
     const std::string &name,
     unsigned int paramCount);
 
-/** Creates RN Runtime Remote Function. */
-jsi::Value makeSerializableRemoteFunction(
+jsi::Value makeRNOriginSerializableRemoteFunction(
     jsi::Runtime &rnRuntime,
     const std::string &name,
     int remoteId,
     const std::shared_ptr<JSScheduler> &jsScheduler);
 
-/** Creates Worklet Runtime Remote Function. */
-jsi::Value makeSerializableRemoteFunction(
+jsi::Value makeWorkletOriginSerializableRemoteFunction(
     jsi::Runtime &workletRuntime,
     const std::string &name,
     jsi::Function &&function,

packages/react-native-worklets/Common/cpp/worklets/SharedItems/SerializableRemoteFunction.cpp

@@ -0,0 +1,113 @@
+#include <jsi/jsi.h>
+#include <react/debug/react_native_assert.h>
+#include <worklets/SharedItems/Serializable.h>
+#include <worklets/SharedItems/SerializableRemoteFunction.h>
+#include <worklets/WorkletRuntime/RuntimeManager.h>
+#include <worklets/WorkletRuntime/WorkletRuntime.h>
+
+#include <memory>
+#include <utility>
+
+using namespace facebook;
+
+namespace worklets {
+
+namespace {
+
+jsi::Function getRemoteFunctionUnpacker(jsi::Runtime &rt) {
+  auto remoteFunctionUnpacker = rt.global().getProperty(rt, "__remoteFunctionUnpacker");
+  react_native_assert(remoteFunctionUnpacker.isObject() && "remoteFunctionUnpacker not found");
+  return remoteFunctionUnpacker.asObject(rt).asFunction(rt);
+}
+
+jsi::Object getRemoteFunctionRegistry(jsi::Runtime &rt) {
+  auto registry = rt.global().getProperty(rt, "__remoteFunctionRegistry");
+  react_native_assert(registry.isObject() && "remoteFunctionRegistry not found");
+  return registry.getObject(rt);
+}
+
+} // namespace
+
+SerializableRemoteFunction::~SerializableRemoteFunction() = default;
+
+jsi::Value SerializableRemoteFunction::unpackSelf(jsi::Runtime &rt) {
+  const auto nameValue = name_.empty() ? jsi::Value::undefined() : jsi::String::createFromUtf8(rt, name_);
+  auto holderFunction = getRemoteFunctionUnpacker(rt).call(rt, nameValue).getObject(rt);
+  holderFunction.setNativeState(rt, std::make_shared<SerializableJSRef>(shared_from_this()));
+  return holderFunction;
+}
+
+jsi::Value SerializableRemoteFunction::RNOrigin::toJSValue(jsi::Runtime &rt) {
+  if (&rt == hostRuntime_) {
+    const auto registry = getRemoteFunctionRegistry(rt);
+    return registry.getPropertyAsFunction(rt, "get").callWithThis(rt, registry, jsi::Value(remoteId_));
+  }
+  std::shared_ptr<SerializableRemoteFunction> proxy;
+  {
+    std::lock_guard<std::mutex> lock(proxyMutex_);
+    proxy = proxy_.lock();
+    if (!proxy) {
+      auto self = std::static_pointer_cast<RNOrigin>(shared_from_this());
+      proxy = std::make_shared<RNOriginProxy>(self);
+      proxy_ = proxy;
+    }
+  }
+  return proxy->toJSValue(rt);
+}
+
+void SerializableRemoteFunction::RNOrigin::resolveOrRejectPromise(
+    const std::shared_ptr<Serializable> &resolveValue,
+    const std::shared_ptr<RuntimeManager> & /*runtimeManager*/) {
+  jsScheduler_->scheduleOnJS([resolver = shared_from_this(), resolveValue](jsi::Runtime &rt) {
+    resolver->toJSValue(rt).getObject(rt).getFunction(rt).call(rt, resolveValue->toJSValue(rt));
+  });
+}
+
+SerializableRemoteFunction::WorkletOrigin::~WorkletOrigin() {
+  cleanupRuntimeAware(hostRuntime_, function_);
+}
+
+jsi::Value SerializableRemoteFunction::WorkletOrigin::toJSValue(jsi::Runtime &rt) {
+  if (&rt == hostRuntime_) {
+    return jsi::Value(rt, *function_);
+  }
+  return unpackSelf(rt);
+}
+
+void SerializableRemoteFunction::WorkletOrigin::resolveOrRejectPromise(
+    const std::shared_ptr<Serializable> &resolveValue,
+    const std::shared_ptr<RuntimeManager> &runtimeManager) {
+  const auto workletRuntime = runtimeManager->getRuntime(hostRuntimeId_);
+  // NOLINTNEXTLINE(readability/braces)
+  if (!workletRuntime) [[unlikely]] {
+    return;
+  }
+  workletRuntime->schedule([resolver = shared_from_this(), resolveValue](jsi::Runtime &rt) {
+    resolver->toJSValue(rt).getObject(rt).getFunction(rt).call(rt, resolveValue->toJSValue(rt));
+  });
+}
+
+SerializableRemoteFunction::RNOrigin::RNOriginProxy::RNOriginProxy(const std::shared_ptr<RNOrigin> &origin)
+    : SerializableRemoteFunction(nullptr, origin->getHostRuntimeId(), origin->getName()), origin_(origin) {}
+
+SerializableRemoteFunction::RNOrigin::RNOriginProxy::~RNOriginProxy() {
+  origin_->getJSScheduler()->scheduleOnJS([id = origin_->getRemoteId()](jsi::Runtime &rt) {
+    const auto registry = getRemoteFunctionRegistry(rt);
+    registry.getPropertyAsFunction(rt, "delete").callWithThis(rt, registry, jsi::Value(id));
+  });
+}
+
+jsi::Value SerializableRemoteFunction::RNOrigin::RNOriginProxy::toJSValue(jsi::Runtime &rt) {
+  if (&rt == origin_->getHostRuntime()) {
+    return origin_->toJSValue(rt);
+  }
+  return unpackSelf(rt);
+}
+
+void SerializableRemoteFunction::RNOrigin::RNOriginProxy::resolveOrRejectPromise(
+    const std::shared_ptr<Serializable> &resolveValue,
+    const std::shared_ptr<RuntimeManager> &runtimeManager) {
+  origin_->resolveOrRejectPromise(resolveValue, runtimeManager);
+}
+
+} // namespace worklets