refactor: explicit invalidations for native and cpp (#6850)

## Summary

This pull requests refactors memory management of Worklets and
Reanimated.

Basically, since Reanimated can obtain WorkletsModuleProxy and the
Worklet Runtimes as shared pointers, it has to release them explicitly
during the invalidation stage of Native Modules. Releasing them later on
(e.g. via deconstructors) might lead into issues and crashes.

Ideally we'd instead use some different solution here than shared
pointers, but it can wait as it's not mandatory at the moment and could
be a significant refactor.

Fixes:
- iOS crash on reload
- Android crash on SingleInstanceChecker during third reload

## Test plan

- [x] 0.76 iOS/Android Fabric works
- [x] 0.76 iOS/Android Paper works
- [x] 0.75 iOS/Android Fabric works
- [x] 0.75 iOS/Android Paper works
- [x] 0.74 iOS/Android Fabric works
- [x] 0.74 iOS/Android Paper works

packages/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxy.cpp

@@ -843,4 +843,10 @@ void ReanimatedModuleProxy::unsubscribeFromKeyboardEvents(
   unsubscribeFromKeyboardEventsFunction_(listenerId.asNumber());
 }
 
+void ReanimatedModuleProxy::invalidate() {
+  // Make sure to release WorkletsModuleProxy on invalidate to allow it
+  // to destroy its runtime during the invalidation stage.
+  workletsModuleProxy_.reset();
+}
+
 } // namespace reanimated

packages/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxy.h

@@ -42,6 +42,8 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
 
   ~ReanimatedModuleProxy();
 
+  void invalidate();
+
   jsi::Value registerEventHandler(
       jsi::Runtime &rt,
       const jsi::Value &worklet,
@@ -175,7 +177,7 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
 
   const bool isBridgeless_;
   const bool isReducedMotion_;
-  const std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
+  std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
   const std::string valueUnpackerCode_;
 
   std::unique_ptr<EventHandlerRegistry> eventHandlerRegistry_;

packages/react-native-reanimated/Common/cpp/worklets/Tools/SingleInstanceChecker.h

@@ -52,10 +52,8 @@ SingleInstanceChecker<T>::SingleInstanceChecker() {
   std::string className =
       __cxxabiv1::__cxa_demangle(typeid(T).name(), nullptr, nullptr, &status);
 
-  // Only one instance should exist, but it is possible for two instances
-  // to co-exist during a reload.
   assertWithMessage(
-      instanceCount_ <= 1,
+      instanceCount_ == 0,
       "[Reanimated] More than one instance of " + className +
           " present. This may indicate a memory leak due to a retain cycle.");
 

packages/react-native-reanimated/android/src/fabric/java/com/swmansion/reanimated/NativeProxy.java

@@ -70,6 +70,12 @@ protected HybridData getHybridData() {
     return mHybridData;
   }
 
+  private native void invalidateCpp();
+
+  public void invalidate() {
+    invalidateCpp();
+  }
+
   public static NativeMethodsHolder createNativeMethodsHolder(
       LayoutAnimations ignoredLayoutAnimations) {
     return new NativeMethodsHolder() {

packages/react-native-reanimated/android/src/main/cpp/reanimated/android/NativeProxy.cpp

@@ -194,7 +194,8 @@ void NativeProxy::registerNatives() {
        makeNativeMethod(
            "isAnyHandlerWaitingForEvent",
            NativeProxy::isAnyHandlerWaitingForEvent),
-       makeNativeMethod("performOperations", NativeProxy::performOperations)});
+       makeNativeMethod("performOperations", NativeProxy::performOperations),
+       makeNativeMethod("invalidateCpp", NativeProxy::invalidateCpp)});
 }
 
 void NativeProxy::requestRender(
@@ -619,4 +620,12 @@ void NativeProxy::setupLayoutAnimations() {
       });
 }
 
+void NativeProxy::invalidateCpp() {
+  workletsModuleProxy_.reset();
+  if (reanimatedModuleProxy_ != nullptr) {
+    reanimatedModuleProxy_->invalidate();
+  }
+  reanimatedModuleProxy_.reset();
+}
+
 } // namespace reanimated

packages/react-native-reanimated/android/src/main/cpp/reanimated/android/NativeProxy.h

@@ -279,6 +279,8 @@ class NativeProxy : public jni::HybridClass<NativeProxy> {
   void commonInit(jni::alias_ref<facebook::react::JFabricUIManager::javaobject>
                       &fabricUIManager);
 #endif // RCT_NEW_ARCH_ENABLED
+
+  void invalidateCpp();
 };
 
 } // namespace reanimated

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

@@ -18,15 +18,13 @@ using namespace facebook;
 using namespace react;
 
 WorkletsModule::WorkletsModule(
-    jni::alias_ref<WorkletsModule::javaobject> jThis,
     jsi::Runtime *rnRuntime,
     const std::string &valueUnpackerCode,
     jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
     const std::shared_ptr<facebook::react::CallInvoker> &jsCallInvoker,
     const std::shared_ptr<worklets::JSScheduler> &jsScheduler,
     const std::shared_ptr<UIScheduler> &uiScheduler)
-    : javaPart_(jni::make_global(jThis)),
-      rnRuntime_(rnRuntime),
+    : rnRuntime_(rnRuntime),
       workletsModuleProxy_(std::make_shared<WorkletsModuleProxy>(
           *rnRuntime,
           valueUnpackerCode,
@@ -38,7 +36,7 @@ WorkletsModule::WorkletsModule(
 }
 
 jni::local_ref<WorkletsModule::jhybriddata> WorkletsModule::initHybrid(
-    jni::alias_ref<jhybridobject> jThis,
+    jni::alias_ref<jhybridobject> /*jThis*/,
     jlong jsContext,
     const std::string &valueUnpackerCode,
     jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
@@ -52,7 +50,6 @@ jni::local_ref<WorkletsModule::jhybriddata> WorkletsModule::initHybrid(
       std::make_shared<worklets::JSScheduler>(*rnRuntime, jsCallInvoker);
   auto uiScheduler = androidUIScheduler->cthis()->getUIScheduler();
   return makeCxxInstance(
-      jThis,
       rnRuntime,
       valueUnpackerCode,
       messageQueueThread,
@@ -61,8 +58,14 @@ jni::local_ref<WorkletsModule::jhybriddata> WorkletsModule::initHybrid(
       uiScheduler);
 }
 
+void WorkletsModule::invalidateCpp() {
+  workletsModuleProxy_.reset();
+}
+
 void WorkletsModule::registerNatives() {
-  registerHybrid({makeNativeMethod("initHybrid", WorkletsModule::initHybrid)});
+  registerHybrid(
+      {makeNativeMethod("initHybrid", WorkletsModule::initHybrid),
+       makeNativeMethod("invalidateCpp", WorkletsModule::invalidateCpp)});
 }
 
 } // namespace worklets

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

@@ -30,7 +30,7 @@ class WorkletsModule : public jni::HybridClass<WorkletsModule> {
       "Lcom/swmansion/worklets/WorkletsModule;";
 
   static jni::local_ref<jhybriddata> initHybrid(
-      jni::alias_ref<jhybridobject> jThis,
+      jni::alias_ref<jhybridobject> /*jThis*/,
       jlong jsContext,
       const std::string &valueUnpackerCode,
       jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
@@ -46,19 +46,19 @@ class WorkletsModule : public jni::HybridClass<WorkletsModule> {
   }
 
  private:
-  friend HybridBase;
-  jni::global_ref<WorkletsModule::javaobject> javaPart_;
-  jsi::Runtime *rnRuntime_;
-  std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
-
   explicit WorkletsModule(
-      jni::alias_ref<WorkletsModule::jhybridobject> jThis,
       jsi::Runtime *rnRuntime,
       const std::string &valueUnpackerCode,
       jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
       const std::shared_ptr<facebook::react::CallInvoker> &jsCallInvoker,
       const std::shared_ptr<worklets::JSScheduler> &jsScheduler,
       const std::shared_ptr<UIScheduler> &uiScheduler);
+
+  void invalidateCpp();
+
+  friend HybridBase;
+  jsi::Runtime *rnRuntime_;
+  std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
 };
 
 } // namespace worklets

packages/react-native-reanimated/android/src/main/java/com/swmansion/reanimated/NodesManager.java

@@ -128,6 +128,7 @@ public void invalidate() {
     }
 
     if (mNativeProxy != null) {
+      mNativeProxy.invalidate();
       mNativeProxy = null;
     }
   }

packages/react-native-reanimated/android/src/main/java/com/swmansion/worklets/WorkletsModule.java

@@ -69,6 +69,13 @@ public boolean installTurboModule(String valueUnpackerCode) {
   }
 
   public void invalidate() {
+    // 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();
   }
+
+  private native void invalidateCpp();
 }

packages/react-native-reanimated/android/src/paper/java/com/swmansion/reanimated/NativeProxy.java

@@ -61,6 +61,12 @@ protected HybridData getHybridData() {
     return mHybridData;
   }
 
+  private native void invalidateCpp();
+
+  public void invalidate() {
+    invalidateCpp();
+  }
+
   public static NativeMethodsHolder createNativeMethodsHolder(LayoutAnimations layoutAnimations) {
     WeakReference<LayoutAnimations> weakLayoutAnimations = new WeakReference<>(layoutAnimations);
     return new NativeMethodsHolder() {

packages/react-native-reanimated/apple/worklets/apple/WorkletsModule.mm

@@ -52,4 +52,14 @@ @implementation WorkletsModule {
   return @YES;
 }
 
+- (void)invalidate
+{
+  // 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 destruction.
+  workletsModuleProxy_.reset();
+
+  [super invalidate];
+}
+
 @end