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`.

(cherry picked from commit 325a644)

Author: tjzel

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

@@ -336,25 +336,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 makeSerializableRemoteFunction(rt, name, fun, jsScheduler, rnRuntimeStatus);
+        } else {
+          return makeSerializableRemoteFunction(rt, name, std::move(fun), hostRuntimeId);
         }
-        if (hostRuntimeId == RuntimeData::rnRuntimeId) {
-          const int remoteId = static_cast<int>(at<1>(args).getNumber());
-          auto ref = makeSerializableRemoteFunction(rt, name, remoteId, jsScheduler);
-          ref.asObject(rt).setProperty(rt, "__keepAlive", true);
-          return ref;
-        }
-        return makeSerializableRemoteFunction(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.cpp

@@ -28,12 +28,6 @@ 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
 
 jsi::Value makeSerializableClone(
@@ -298,28 +292,19 @@ jsi::Value SerializableImport::toJSValue(jsi::Runtime &rt) {
 
 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));
-    });
+    if (rnRuntimeStatus_->isDead()) {
+      freeWithoutCallingDestructor(function_);
+    } else {
+      function_.reset();
+    }
   } else {
-    auto &workletData = std::get<WorkletRuntimeData>(runtimeData_);
-    cleanupRuntimeAware(hostRuntime_, workletData.function);
+    cleanupRuntimeAware(hostRuntime_, 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);
-    }
+    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);
@@ -333,8 +318,7 @@ 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) {
+    jsScheduler_->scheduleOnJS([resolver = shared_from_this(), resolveValue](jsi::Runtime &rt) {
       resolver->toJSValue(rt).getObject(rt).getFunction(rt).call(rt, resolveValue->toJSValue(rt));
     });
   } else {

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

@@ -4,13 +4,13 @@
 #include <worklets/Compat/StableApi.h>
 #include <worklets/Registries/WorkletRuntimeRegistry.h>
 #include <worklets/Tools/JSScheduler.h>
+#include <worklets/Tools/RNRuntimeStatus.h>
 #include <worklets/WorkletRuntime/RuntimeData.h>
 
 #include <memory>
 #include <optional>
 #include <string>
 #include <utility>
-#include <variant>
 #include <vector>
 
 using namespace facebook;
@@ -283,33 +283,21 @@ 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)
+      jsi::Function &&function,
+      const std::shared_ptr<JSScheduler> &jsScheduler,
+      const std::shared_ptr<RNRuntimeStatus> &rnRuntimeStatus)
       : Serializable(ValueType::RemoteFunctionType),
         hostRuntime_(&rnRuntime),
         hostRuntimeId_(RuntimeData::rnRuntimeId),
-        runtimeData_(RNRuntimeData{.remoteId = remoteId, .jsScheduler = jsScheduler}),
-        name_(name) {}
+        function_(std::make_unique<jsi::Value>(rnRuntime, std::move(function))),
+        name_(name),
+        jsScheduler_(jsScheduler),
+        rnRuntimeStatus_(rnRuntimeStatus) {}
 
   /** Creates Worklet Runtime Remote Function. */
   SerializableRemoteFunction(
@@ -320,8 +308,10 @@ class SerializableRemoteFunction : public Serializable,
       : Serializable(ValueType::RemoteFunctionType),
         hostRuntime_(&workletRuntime),
         hostRuntimeId_(hostRuntimeId),
-        runtimeData_(WorkletRuntimeData{.function = std::make_unique<jsi::Value>(workletRuntime, std::move(function))}),
-        name_(name) {}
+        function_(std::make_unique<jsi::Value>(workletRuntime, std::move(function))),
+        name_(name),
+        jsScheduler_(nullptr),
+        rnRuntimeStatus_(nullptr) {}
 
   ~SerializableRemoteFunction() override;
 
@@ -335,12 +325,20 @@ class SerializableRemoteFunction : public Serializable,
   jsi::Value toJSValue(jsi::Runtime &rt) override;
 
   [[nodiscard]] bool isHostedOnRNRuntime() const noexcept {
-    return std::holds_alternative<RNRuntimeData>(runtimeData_);
+    return hostRuntimeId_ == RuntimeData::rnRuntimeId;
   }
 
   [[nodiscard]] RuntimeData::RuntimeId getHostRuntimeId() const {
     return hostRuntimeId_;
   }
+
+ private:
+  jsi::Runtime *hostRuntime_;
+  const RuntimeData::RuntimeId hostRuntimeId_;
+  std::unique_ptr<jsi::Value> function_;
+  const std::string name_;
+  const std::shared_ptr<JSScheduler> jsScheduler_;
+  const std::shared_ptr<RNRuntimeStatus> rnRuntimeStatus_;
 };
 
 class SerializableInitializer : public Serializable {

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

@@ -63,9 +63,15 @@ jsi::Value makeSerializableHostFunction(
 jsi::Value makeSerializableRemoteFunction(
     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);
+    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);
 }
 

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

@@ -63,8 +63,9 @@ jsi::Value makeSerializableHostFunction(
 jsi::Value makeSerializableRemoteFunction(
     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);
 
 /** Creates Worklet Runtime Remote Function. */
 jsi::Value makeSerializableRemoteFunction(

packages/react-native-worklets/Common/cpp/worklets/Tools/RNRuntimeStatus.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <atomic>
+
+namespace worklets {
+
+class RNRuntimeStatus {
+ public:
+  [[nodiscard]] bool isDead() const {
+    return isDead_;
+  }
+
+  void setDead() {
+    isDead_ = true;
+  }
+
+ private:
+  std::atomic_bool isDead_ = false;
+};
+
+} // namespace worklets

packages/react-native-worklets/android/src/main/cpp/worklets/android/WorkletsModule.cpp

@@ -33,13 +33,15 @@ WorkletsModule::WorkletsModule(
     const std::shared_ptr<UIScheduler> &uiScheduler)
     : javaPart_(jni::make_global(jThis)),
       rnRuntime_(rnRuntime),
+      rnRuntimeStatus_(std::make_shared<RNRuntimeStatus>()),
       workletsModuleProxy_(std::make_shared<WorkletsModuleProxy>(
           *rnRuntime,
           jsCallInvoker,
           uiScheduler,
           getIsOnJSQueueThread(),
           getRuntimeBindings(bundleModeConfig.enabled, *rnRuntime),
-          bundleModeConfig)) {}
+          bundleModeConfig,
+          rnRuntimeStatus_)) {}
 
 jni::local_ref<WorkletsModule::jhybriddata> WorkletsModule::initHybrid(
     jni::alias_ref<jhybridobject> jThis, // NOLINT //(performance-unnecessary-value-param)
@@ -179,6 +181,7 @@ std::function<bool()> WorkletsModule::getIsOnJSQueueThread() {
 }
 
 void WorkletsModule::invalidateCpp() {
+  rnRuntimeStatus_->setDead();
   javaPart_.reset();
   workletsModuleProxy_.reset();
 }