fix(Worklets): RemoteFunction memory cycle (#9673)

## Summary

Fixes #9661

In #9272 which was supposed to fix memory issues with Remote Function, I
accidentally introduced a memory-leak due to reference retain cycle.

### Before

Basically, RemoteFunctionRegistry would hold a strong reference to the
(remote) function, which in turn would hold a strong reference to the
serialized remote function, which would hold the C++ value.

This C++ value would erase the original function from the registry, but
only in its destructor - which would never happen, because the registry
itself would keep it alive forever.

```mermaid
flowchart TB
  subgraph RN[RN Runtime]
    R[RemoteFunctionRegistry]
    F[remote function]
    C[serialized remote function]
  end
  subgraph CPP[C++ heap]
    S[C++ value]
  end
  subgraph WK[Worker Runtime]
    H[holderFunction]
  end

  R --> F
  F -.-> C
  C --> S
  H --> S
  S ==>|"registry.delete(id) — never invoked"| R

  linkStyle 0,2,3 stroke:#2f9e44,stroke-width:2px
  linkStyle 1 stroke:#868e96,stroke-width:2px
  linkStyle 4 stroke:#e03131,stroke-width:3px
```

### After

To fix this issue, I split the C++ Remote Function for RN callbacks into
two subclasses - one would only be owned by the RN Runtime, and the
other one from Worklet Runtimes.

If all Worklet Runtime references would get cleaned up, then they would
schedule cleanup on the RN Runtime C++ side, to erase the original
function from the registry. Then, this would in turn release the cache
value and the remaining C++ value.


```mermaid
flowchart TB
  subgraph RN[RN Runtime]
    R[RemoteFunctionRegistry]
    F[remote function]
    C[serialized remote function]
  end
  subgraph CPP[C++ heap]
    O[RN-owned C++ value]
    P[Worklet-owned C++ value]
  end
  subgraph WK[Worker Runtime]
    H[holderFunction]
  end

  R --> F
  F -.-> C
  C --> O
  H --> P
  P --> O
  O -.-> P
  P ==>|"registry.delete(id)"| R

  linkStyle 0,2,3,4 stroke:#2f9e44,stroke-width:2px
  linkStyle 1,5 stroke:#868e96,stroke-width:2px
  linkStyle 6 stroke:#2f9e44,stroke-width:3px
```

## Test Plan

I used the following code and was tracking constructor and destructor
invocations of `SerializableRemoteFunction` in debugger.

<details><summary>Code</summary>
<p>

```tsx
import { useEffect } from 'react';
import { Button, StyleSheet, View } from 'react-native';
import {
  scheduleOnRN,
  scheduleOnUI,
  createWorkletRuntime,
  scheduleOnRuntime,
  createSerializable,
} from 'react-native-worklets';

const workerRuntime = createWorkletRuntime();

export default function App() {
  useEffect(() => {
    scheduleOnUI(() => {
      function foo() {
        const _ = new Array(10000).fill(0);
        requestAnimationFrame(foo);
      }
      foo();
    });
  });

  return (
    <View style={styles.container}>
      <Button
        title="Init runtime"
        onPress={() => {
          scheduleOnUI(() => {
            // scheduleOnRuntime(workerRuntime, () => {
            'worklet';
            console.log('init');
          });
        }}
      />
      <Button
        title="Press me first"
        onPress={() => {
          function logTest() {
            console.log('Button pressed!');
          }

          // scheduleOnRuntime(workerRuntime, () => {
          scheduleOnUI(() => {
            'worklet';
            scheduleOnRN(logTest);
          });
        }}
      />
      <Button
        title="Trigger GC on RN"
        onPress={() => {
          console.log('Triggering GC on RN');
          gc();
        }}
      />
      <Button
        title="Trigger GC on Worker"
        onPress={() => {
          // scheduleOnRuntime(workerRuntime, () => {
          scheduleOnUI(() => {
            'worklet';
            console.log('Triggering GC on Worker');
            gc();
          });
        }}
      />
    </View>
  );
}

const styles = StyleSheet.create({
  container: {
    flex: 1,
    backgroundColor: '#fff',
    alignItems: 'center',
    justifyContent: 'center',
  },
});
```

</p>
</details>

Author: tjzel

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(
@@ -302,61 +291,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>
@@ -289,71 +290,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

@@ -63,15 +63,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