cherry-pick(worklets-0.9-stable): tie RNR remote function lifetime to TurboModule invalidate status (#9789)

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();
 }

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

@@ -4,6 +4,7 @@
 #include <fbjni/fbjni.h>
 #include <jsi/jsi.h>
 #include <worklets/NativeModules/WorkletsModuleProxy.h>
+#include <worklets/Tools/RNRuntimeStatus.h>
 #include <worklets/Tools/ScriptBuffer.h>
 #include <worklets/WorkletRuntime/BundleModeConfig.h>
 #include <worklets/WorkletRuntime/RuntimeBindings.h>
@@ -67,6 +68,7 @@ class WorkletsModule : public jni::HybridClass<WorkletsModule> {
   friend HybridBase;
   jni::global_ref<WorkletsModule::javaobject> javaPart_;
   jsi::Runtime *rnRuntime_;
+  std::shared_ptr<RNRuntimeStatus> rnRuntimeStatus_;
   std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
 };
 

packages/react-native-worklets/android/src/networking/com/swmansion/worklets/WorkletsModule.kt

@@ -166,9 +166,6 @@ class WorkletsModule(
             return
         }
         if (mHybridData != null && mHybridData!!.isValid) {
-            // We have to destroy extra runtimes when invalidate is called. If we clean
-            // it up later instead there's a chance the runtime will retain references
-            // to invalidated memory and will crash on its destruction.
             invalidateCpp()
         }
         mAndroidUIScheduler.deactivate()

packages/react-native-worklets/android/src/no-networking/com/swmansion/worklets/WorkletsModule.kt

@@ -122,9 +122,6 @@ class WorkletsModule(
             return
         }
         if (mHybridData != null && mHybridData!!.isValid) {
-            // We have to destroy extra runtimes when invalidate is called. If we clean
-            // it up later instead there's a chance the runtime will retain references
-            // to invalidated memory and will crash on its destruction.
             invalidateCpp()
         }
         mAndroidUIScheduler.deactivate()