refactor: explicit invalidations for native and cpp

Author: tjzel

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

@@ -61,8 +61,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

@@ -46,11 +46,6 @@ 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,
@@ -59,6 +54,13 @@ class WorkletsModule : public jni::HybridClass<WorkletsModule> {
       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;
+  jni::global_ref<WorkletsModule::javaobject> javaPart_;
+  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 destruction.
+    invalidateCpp();
     mAndroidUIScheduler.deactivate();
   }
+
+  private native void invalidateCpp();
 }