fix(Worklets): tie RNR remote function lifetime to TurboModule invalidate status (#9789)

## Summary

Fixes #9751 

Changing (again) how RN Runtime remote functions (RFs) are managed,
since changes from #9673 and #9272 were actually insufficient and the
map approach won't work.

It's impossible to keep both the actual callback and cached serialized
callback value alive without a memory cycle, so I'm dropping cache
consistency. This might result in multiple serializations of a given RF
the GC triggered and the serialized value wasn't saved, but the
serialization of a RF is actually cheap.

I'm dropping the map approach entirely in favor of the behavior where
destructor invocation is based upon the information if the runtime is
still alive. We can't directly check it, so I'm depending on the
`invalidate()` method of Turbo Modules - which is supposed to be called
when the runtime is still alive.

## Test plan

You can test it on
[df3e344](df3e344)
commit - there's a special registry there that artificially extends the
lifetime of Remote Functions - click some buttons on the provided
example, then reload the app with R button. The callbacks with persist
and will crash on destructor invocation without the check for `isDead`.

Author: tjzel

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

@@ -335,25 +335,22 @@ jsi::Object JSIWorkletsModuleProxy::toOptimizedObject(jsi::Runtime &rt) const {
         return makeSerializableArrayBuffer(rt, buffer);
       });
 
-  jsi_utils::addMethod<3>(
+  jsi_utils::addMethod<2>(
       rt,
       obj,
       "createSerializableNonWorkletFunction",
-      [jsScheduler = jsScheduler_, hostRuntimeId = hostRuntimeId_](
-          jsi::Runtime &rt, const jsi::Value &, const jsi::Value(&args)[3]) {
+      [jsScheduler = jsScheduler_, hostRuntimeId = hostRuntimeId_, rnRuntimeStatus = rnRuntimeStatus_](
+          jsi::Runtime &rt, const jsi::Value &, const jsi::Value(&args)[2]) {
         auto fun = at<0>(args).getObject(rt).getFunction(rt);
-        const auto name = at<2>(args).isUndefined() ? "" : at<2>(args).getString(rt).utf8(rt);
+        const auto name = at<1>(args).isUndefined() ? "" : at<1>(args).getString(rt).utf8(rt);
         if (fun.isHostFunction(rt)) {
           return makeSerializableHostFunction(
               rt, fun.getHostFunction(rt), name, fun.getProperty(rt, "length").getNumber());
+        } else if (hostRuntimeId == RuntimeData::rnRuntimeId) {
+          return makeRNRuntimeSerializableRemoteFunction(rt, name, fun, jsScheduler, rnRuntimeStatus);
+        } else {
+          return makeWorkletRuntimeSerializableRemoteFunction(rt, name, fun, hostRuntimeId);
         }
-        if (hostRuntimeId == RuntimeData::rnRuntimeId) {
-          const int remoteId = static_cast<int>(at<1>(args).getNumber());
-          auto ref = makeRNOriginSerializableRemoteFunction(rt, name, remoteId, jsScheduler);
-          ref.asObject(rt).setProperty(rt, "__keepAlive", true);
-          return ref;
-        }
-        return makeWorkletOriginSerializableRemoteFunction(rt, name, std::move(fun), hostRuntimeId);
       });
 
   jsi_utils::addMethod<2>(

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

@@ -4,6 +4,7 @@
 #include <worklets/SharedItems/MemoryManager.h>
 #include <worklets/SharedItems/Serializable.h>
 #include <worklets/SharedItems/UnpackerLoader.h>
+#include <worklets/Tools/RNRuntimeStatus.h>
 #include <worklets/Tools/ScriptBuffer.h>
 #include <worklets/WorkletRuntime/BundleModeConfig.h>
 #include <worklets/WorkletRuntime/RuntimeBindings.h>
@@ -31,6 +32,7 @@ class JSIWorkletsModuleProxy : public std::enable_shared_from_this<JSIWorkletsMo
       const std::shared_ptr<RuntimeBindings> &runtimeBindings,
       const BundleModeConfig &bundleModeConfig,
       const std::shared_ptr<UnpackerLoader> &unpackerLoader,
+      const std::shared_ptr<RNRuntimeStatus> &rnRuntimeStatus,
       RuntimeData::RuntimeId hostRuntimeId)
       : isDevBundle_(isDevBundle),
         bundleModeConfig_(bundleModeConfig),
@@ -41,6 +43,7 @@ class JSIWorkletsModuleProxy : public std::enable_shared_from_this<JSIWorkletsMo
         uiWorkletRuntime_(uiWorkletRuntime),
         runtimeBindings_(runtimeBindings),
         unpackerLoader_(unpackerLoader),
+        rnRuntimeStatus_(rnRuntimeStatus),
         hostRuntimeId_(hostRuntimeId) {}
 
   static std::shared_ptr<JSIWorkletsModuleProxy> createForNewRuntime(
@@ -56,6 +59,7 @@ class JSIWorkletsModuleProxy : public std::enable_shared_from_this<JSIWorkletsMo
         sourceProxy->runtimeBindings_,
         sourceProxy->bundleModeConfig_,
         sourceProxy->unpackerLoader_,
+        sourceProxy->rnRuntimeStatus_,
         hostRuntimeId);
   }
 
@@ -115,6 +119,7 @@ class JSIWorkletsModuleProxy : public std::enable_shared_from_this<JSIWorkletsMo
   const std::weak_ptr<WorkletRuntime> uiWorkletRuntime_;
   const std::shared_ptr<RuntimeBindings> runtimeBindings_;
   const std::shared_ptr<UnpackerLoader> unpackerLoader_;
+  const std::shared_ptr<RNRuntimeStatus> rnRuntimeStatus_;
   const RuntimeData::RuntimeId hostRuntimeId_;
 };
 

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

@@ -45,7 +45,8 @@ WorkletsModuleProxy::WorkletsModuleProxy(
     const std::shared_ptr<UIScheduler> &uiScheduler,
     std::function<bool()> &&isJavaScriptThread,
     const std::shared_ptr<RuntimeBindings> &runtimeBindings,
-    const BundleModeConfig &bundleModeConfig)
+    const BundleModeConfig &bundleModeConfig,
+    const std::shared_ptr<RNRuntimeStatus> &rnRuntimeStatus)
     : isDevBundle_(isDevBundleFromRNRuntime(rnRuntime)),
       jsScheduler_(std::make_shared<JSScheduler>(rnRuntime, jsCallInvoker, std::move(isJavaScriptThread))),
       uiScheduler_(uiScheduler),
@@ -55,6 +56,7 @@ WorkletsModuleProxy::WorkletsModuleProxy(
       memoryManager_(std::make_shared<MemoryManager>()),
       runtimeManager_(std::make_shared<RuntimeManager>()),
       unpackerLoader_(std::make_shared<UnpackerLoader>()),
+      rnRuntimeStatus_(rnRuntimeStatus),
       uiWorkletRuntime_(runtimeManager_->createUninitializedUIRuntime(std::make_shared<AsyncQueueUI>(uiScheduler_))),
       rnRuntimeProxy_(std::make_shared<JSIWorkletsModuleProxy>(
           isDevBundle_,
@@ -66,6 +68,7 @@ WorkletsModuleProxy::WorkletsModuleProxy(
           runtimeBindings_,
           bundleModeConfig_,
           unpackerLoader_,
+          rnRuntimeStatus_,
           RuntimeData::rnRuntimeId)) {
   RNRuntimeWorkletDecorator::decorate(rnRuntime, rnRuntimeProxy_->toOptimizedObject(rnRuntime), jsLogger_);
 }

packages/react-native-worklets/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.h

@@ -7,6 +7,7 @@
 #include <worklets/SharedItems/UnpackerLoader.h>
 #include <worklets/Tools/JSLogger.h>
 #include <worklets/Tools/JSScheduler.h>
+#include <worklets/Tools/RNRuntimeStatus.h>
 #include <worklets/Tools/ScriptBuffer.h>
 #include <worklets/Tools/SingleInstanceChecker.h>
 #include <worklets/Tools/UIScheduler.h>
@@ -28,7 +29,8 @@ class WorkletsModuleProxy : public std::enable_shared_from_this<WorkletsModulePr
       const std::shared_ptr<UIScheduler> &uiScheduler,
       std::function<bool()> &&isJavaScriptQueue,
       const std::shared_ptr<RuntimeBindings> &runtimeBindings,
-      const BundleModeConfig &bundleModeConfig);
+      const BundleModeConfig &bundleModeConfig,
+      const std::shared_ptr<RNRuntimeStatus> &rnRuntimeStatus);
 
   ~WorkletsModuleProxy();
 
@@ -62,6 +64,7 @@ class WorkletsModuleProxy : public std::enable_shared_from_this<WorkletsModulePr
   const std::shared_ptr<MemoryManager> memoryManager_;
   const std::shared_ptr<RuntimeManager> runtimeManager_;
   const std::shared_ptr<UnpackerLoader> unpackerLoader_;
+  const std::shared_ptr<RNRuntimeStatus> rnRuntimeStatus_;
   std::shared_ptr<WorkletRuntime> uiWorkletRuntime_;
   const std::shared_ptr<JSIWorkletsModuleProxy> rnRuntimeProxy_;
   std::shared_ptr<AnimationFrameBatchinator> animationFrameBatchinator_;

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

@@ -11,7 +11,6 @@
 #include <optional>
 #include <string>
 #include <utility>
-#include <variant>
 #include <vector>
 
 using namespace facebook;

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

@@ -61,22 +61,31 @@ jsi::Value makeSerializableHostFunction(
   return SerializableJSRef::newNativeStateObject(rt, serializable);
 }
 
-jsi::Value makeRNOriginSerializableRemoteFunction(
+jsi::Value makeRNRuntimeSerializableRemoteFunction(
     jsi::Runtime &rnRuntime,
     const std::string &name,
-    const int remoteId,
-    const std::shared_ptr<JSScheduler> &jsScheduler) {
-  auto serializable = std::make_shared<SerializableRemoteFunction::RNOrigin>(rnRuntime, name, remoteId, jsScheduler);
+    const jsi::Function &function,
+    const std::shared_ptr<JSScheduler> &jsScheduler,
+    const std::shared_ptr<RNRuntimeStatus> &rnRuntimeStatus) {
+  auto serializable = std::make_shared<SerializableRemoteFunction>(
+      rnRuntime,
+      name,
+      jsi::Value(rnRuntime, function).getObject(rnRuntime).getFunction(rnRuntime),
+      jsScheduler,
+      rnRuntimeStatus);
   return SerializableJSRef::newNativeStateObject(rnRuntime, serializable);
 }
 
-jsi::Value makeWorkletOriginSerializableRemoteFunction(
+jsi::Value makeWorkletRuntimeSerializableRemoteFunction(
     jsi::Runtime &workletRuntime,
     const std::string &name,
-    jsi::Function &&function,
+    const jsi::Function &function,
     RuntimeData::RuntimeId hostRuntimeId) {
-  auto serializable = std::make_shared<SerializableRemoteFunction::WorkletOrigin>(
-      workletRuntime, name, std::move(function), hostRuntimeId);
+  auto serializable = std::make_shared<SerializableRemoteFunction>(
+      workletRuntime,
+      name,
+      jsi::Value(workletRuntime, function).getObject(workletRuntime).getFunction(workletRuntime),
+      hostRuntimeId);
   return SerializableJSRef::newNativeStateObject(workletRuntime, serializable);
 }
 

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

@@ -9,6 +9,8 @@
 
 namespace worklets {
 
+class RNRuntimeStatus;
+
 jsi::Value makeSerializableString(jsi::Runtime &rt, const jsi::String &string);
 
 jsi::Value makeSerializableNumber(jsi::Runtime &rt, double number);
@@ -63,16 +65,17 @@ jsi::Value makeSerializableHostFunction(
     const std::string &name,
     unsigned int paramCount);
 
-jsi::Value makeRNOriginSerializableRemoteFunction(
+jsi::Value makeRNRuntimeSerializableRemoteFunction(
     jsi::Runtime &rnRuntime,
     const std::string &name,
-    int remoteId,
-    const std::shared_ptr<JSScheduler> &jsScheduler);
+    const jsi::Function &function,
+    const std::shared_ptr<JSScheduler> &jsScheduler,
+    const std::shared_ptr<RNRuntimeStatus> &rnRuntimeStatus);
 
-jsi::Value makeWorkletOriginSerializableRemoteFunction(
+jsi::Value makeWorkletRuntimeSerializableRemoteFunction(
     jsi::Runtime &workletRuntime,
     const std::string &name,
-    jsi::Function &&function,
+    const jsi::Function &function,
     RuntimeData::RuntimeId hostRuntimeId);
 
 jsi::Value makeSerializableWorklet(jsi::Runtime &rt, const jsi::Object &object, const bool &shouldRetainRemote);

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

@@ -6,7 +6,6 @@
 #include <worklets/WorkletRuntime/WorkletRuntime.h>
 
 #include <memory>
-#include <utility>
 
 using namespace facebook;
 
@@ -20,94 +19,51 @@ jsi::Function getRemoteFunctionUnpacker(jsi::Runtime &rt) {
   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;
+SerializableRemoteFunction::~SerializableRemoteFunction() {
+  if (isHostedOnRNRuntime()) {
+    if (rnRuntimeStatus_->isDead()) {
+      freeWithoutCallingDestructor(function_);
+    } else {
+      function_.reset();
     }
+  } else {
+    cleanupRuntimeAware(hostRuntime_, function_);
   }
-  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) {
+jsi::Value SerializableRemoteFunction::toJSValue(jsi::Runtime &rt) {
   if (&rt == hostRuntime_) {
     return jsi::Value(rt, *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;
   }
-  return unpackSelf(rt);
 }
 
-void SerializableRemoteFunction::WorkletOrigin::resolveOrRejectPromise(
+// TODO: generalize it and merge with other scheduling methods
+void SerializableRemoteFunction::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);
+  if (isHostedOnRNRuntime()) {
+    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));
+      });
+    }
   }
-  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