feat: ui scheduling mutex

software-mansion · Jan 13, 2025 · fa1f058 · fa1f058
1 parent 15b3351
commit fa1f058
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.cpp b/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.cpp
@@ -52,8 +52,15 @@ WorkletsModuleProxy::WorkletsModuleProxy(
           valueUnpackerCode_)) {}
 
 WorkletsModuleProxy::~WorkletsModuleProxy() {
+  // We need an exclusive lock to make sure nothing can be scheduled on the UI
+  // during the destructor invocation.
+  uiWorkletRuntimeMutex_.lock();
+
+  isValid_ = false;
   jsQueue_->quitSynchronous();
   uiWorkletRuntime_.reset();
+
+  uiWorkletRuntimeMutex_.unlock();
 }
 
 jsi::Value WorkletsModuleProxy::makeShareableClone(
@@ -70,6 +77,13 @@ jsi::Value WorkletsModuleProxy::makeShareableClone(
 void WorkletsModuleProxy::scheduleOnUI(
     jsi::Runtime &rt,
     const jsi::Value &worklet) {
+  // We need a shared lock here to make sure the UI runtime is not being
+  // destroyed while we are scheduling worklets on it.
+  uiWorkletRuntimeMutex_.lock_shared();
+  if (!isValid_) {
+    return;
+  }
+
   auto shareableWorklet = extractShareableOrThrow<ShareableWorklet>(
       rt, worklet, "[Worklets] Only worklets can be scheduled to run on UI.");
   uiScheduler_->scheduleOnUI(COPY_CAPTURE_WITH_THIS {
@@ -82,6 +96,8 @@ void WorkletsModuleProxy::scheduleOnUI(
     const auto scope = jsi::Scope(uiWorkletRuntime_->getJSIRuntime());
 #endif
     uiWorkletRuntime_->runGuarded(shareableWorklet);
+
+    uiWorkletRuntimeMutex_.unlock_shared();
   });
 }
 

diff --git a/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.h b/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.h
@@ -7,6 +7,7 @@
 #include <worklets/Tools/UIScheduler.h>
 #include <worklets/WorkletRuntime/WorkletRuntime.h>
 #include <memory>
+#include <shared_mutex>
 #include <string>
 
 namespace worklets {
@@ -74,6 +75,8 @@ class WorkletsModuleProxy : public WorkletsModuleProxySpec {
 #ifndef NDEBUG
   SingleInstanceChecker<WorkletsModuleProxy> singleInstanceChecker_;
 #endif // NDEBUG
+  bool isValid_{true};
+  std::shared_mutex uiWorkletRuntimeMutex_;
 };
 
 } // namespace worklets