fix(Worklets): time-of-check and time-of-use destructor safety (#9790)

## Summary

Currently the way we check if the runtime is alive for safe destructor
invocations has a chance to lead to crashes anyway, since the runtime
can get cleaned up immediately after the check, since it lacks a locking
mechanism.

## Test plan

The test example from #9789 works.

(cherry picked from commit 71774cb)

Author: tjzel

packages/react-native-worklets/Common/cpp/worklets/Registries/WorkletRuntimeRegistry.cpp

@@ -5,6 +5,6 @@
 namespace worklets {
 
 std::set<jsi::Runtime *> WorkletRuntimeRegistry::registry_{};
-std::mutex WorkletRuntimeRegistry::mutex_{};
+std::shared_mutex WorkletRuntimeRegistry::mutex_{};
 
 } // namespace worklets

packages/react-native-worklets/Common/cpp/worklets/Registries/WorkletRuntimeRegistry.h

@@ -3,8 +3,9 @@
 #include <jsi/jsi.h>
 #include <react/debug/react_native_assert.h>
 
-#include <mutex>
 #include <set>
+#include <shared_mutex>
+#include <utility>
 
 using namespace facebook;
 
@@ -13,27 +14,29 @@ namespace worklets {
 class WorkletRuntimeRegistry {
  private:
   static std::set<jsi::Runtime *> registry_;
-  static std::mutex mutex_; // Protects `registry_`.
+  static std::shared_mutex mutex_;
 
-  WorkletRuntimeRegistry() {} // private ctor
+  WorkletRuntimeRegistry() {}
 
   static void registerRuntime(jsi::Runtime &runtime) {
-    std::lock_guard<std::mutex> lock(mutex_);
+    std::lock_guard<std::shared_mutex> lock(mutex_);
     registry_.insert(&runtime);
   }
 
   static void unregisterRuntime(jsi::Runtime &runtime) {
-    std::lock_guard<std::mutex> lock(mutex_);
+    std::lock_guard<std::shared_mutex> lock(mutex_);
     registry_.erase(&runtime);
   }
 
   friend class WorkletRuntimeCollector;
 
  public:
-  static bool isRuntimeAlive(jsi::Runtime *runtime) {
+  template <typename TFn>
+  static void runWhileLocked(jsi::Runtime *runtime, TFn &&fn) {
     react_native_assert(runtime != nullptr && "runtime is nullptr");
-    std::lock_guard<std::mutex> lock(mutex_);
-    return registry_.find(runtime) != registry_.end();
+    std::shared_lock<std::shared_mutex> lock(mutex_);
+    const bool isAlive = registry_.find(runtime) != registry_.end();
+    std::forward<TFn>(fn)(isAlive);
   }
 };
 

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

@@ -30,9 +30,16 @@ inline void freeWithoutCallingDestructor(std::unique_ptr<jsi::Value> &value) {
 }
 
 inline void cleanupRuntimeAware(jsi::Runtime *rt, std::unique_ptr<jsi::Value> &value) {
-  if (rt != nullptr && !WorkletRuntimeRegistry::isRuntimeAlive(rt)) {
-    freeWithoutCallingDestructor(value);
+  if (value == nullptr || rt == nullptr) {
+    return;
   }
+  WorkletRuntimeRegistry::runWhileLocked(rt, [&value](bool isAlive) {
+    if (isAlive) {
+      value.reset();
+    } else {
+      freeWithoutCallingDestructor(value);
+    }
+  });
 }
 
 template <typename BaseClass>

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

@@ -23,11 +23,13 @@ jsi::Function getRemoteFunctionUnpacker(jsi::Runtime &rt) {
 
 SerializableRemoteFunction::~SerializableRemoteFunction() {
   if (isHostedOnRNRuntime()) {
-    if (rnRuntimeStatus_->isDead()) {
-      freeWithoutCallingDestructor(function_);
-    } else {
-      function_.reset();
-    }
+    rnRuntimeStatus_->runWhileLocked([this](bool isDead) {
+      if (isDead) {
+        freeWithoutCallingDestructor(function_);
+      } else {
+        function_.reset();
+      }
+    });
   } else {
     cleanupRuntimeAware(hostRuntime_, function_);
   }

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

@@ -1,21 +1,27 @@
 #pragma once
 
-#include <atomic>
+#include <shared_mutex>
+#include <utility>
 
 namespace worklets {
 
 class RNRuntimeStatus {
  public:
-  [[nodiscard]] bool isDead() const {
-    return isDead_;
-  }
-
   void setDead() {
+    std::lock_guard<std::shared_mutex> lock(mutex_);
     isDead_ = true;
   }
 
+  template <typename TFn>
+  void runWhileLocked(TFn &&fn) {
+    std::shared_lock<std::shared_mutex> lock(mutex_);
+    const bool isDead = isDead_;
+    std::forward<TFn>(fn)(isDead);
+  }
+
  private:
-  std::atomic_bool isDead_ = false;
+  bool isDead_ = false;
+  std::shared_mutex mutex_;
 };
 
 } // namespace worklets