refactor: explicit invalidations for native and cpp

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

@@ -844,4 +844,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_ < 1,
       "[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,10 @@ void NativeProxy::setupLayoutAnimations() {
       });
 }
 
+void NativeProxy::invalidateCpp() {
+  workletsModuleProxy_.reset();
+  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,12 @@ 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,13 @@ @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