fix: Break circular reference in NativeProxy and AndroidUIScheduler (#6697)

## Summary

Breaks circular reference in `NativeProxy` and `AndroidUIScheduler`.
Needs
#6671 to
work.

`NativeProxy` and `AndroidUIScheduler` is a `HybridClass` with a C++
counterpart holding a global reference to the Java object. This
structure creates a circular reference between C++ and Java, which the
garbage collector cannot clean up. To resolve this issue, I manually
removed the reference between C++ and Java by resetting the global ref
during the invalidation of the `NativeProxy` and `AndroidUIScheduler`.

## Test plan

I tested it with a new project using Expo (SDK 52) and with Fabric
enabled. I reloaded the app several times and performed a heap dump to
verify that everything was removed correctly.

---------

Co-authored-by: Tomasz Żelawski <tzelawski@gmail.com>
Co-authored-by: Krzysztof Piaskowy <krzysztof.piaskowy@swmansion.com>
Co-authored-by: Tomek Zawadzki <tomekzawadzki98@gmail.com>

Author: 4 people

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

@@ -217,6 +217,11 @@ void NativeProxy::registerEventHandler() {
 }
 
 void NativeProxy::maybeFlushUIUpdatesQueue() {
+  // Module might be already destroyed.
+  if (!javaPart_) {
+    return;
+  }
+
   static const auto method = getJniMethod<void()>("maybeFlushUIUpdatesQueue");
   method(javaPart_.get());
 }

packages/react-native-worklets/android/src/main/cpp/worklets/android/AndroidUIScheduler.cpp

@@ -25,8 +25,6 @@ class UISchedulerWrapper : public UIScheduler {
       androidUiScheduler_->cthis()->scheduleTriggerOnUI();
     }
   }
-
-  ~UISchedulerWrapper() {}
 };
 
 AndroidUIScheduler::AndroidUIScheduler(
@@ -41,6 +39,9 @@ jni::local_ref<AndroidUIScheduler::jhybriddata> AndroidUIScheduler::initHybrid(
 }
 
 void AndroidUIScheduler::triggerUI() {
+  if (!uiScheduler_) {
+    return;
+  }
   uiScheduler_->triggerUI();
 }
 
@@ -50,10 +51,16 @@ void AndroidUIScheduler::scheduleTriggerOnUI() {
   method(javaPart_.get());
 }
 
+void AndroidUIScheduler::invalidate() {
+  javaPart_ = nullptr;
+  uiScheduler_.reset();
+}
+
 void AndroidUIScheduler::registerNatives() {
   registerHybrid({
       makeNativeMethod("initHybrid", AndroidUIScheduler::initHybrid),
       makeNativeMethod("triggerUI", AndroidUIScheduler::triggerUI),
+      makeNativeMethod("invalidate", AndroidUIScheduler::invalidate),
   });
 }
 

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

@@ -34,6 +34,8 @@ class AndroidUIScheduler : public jni::HybridClass<AndroidUIScheduler> {
 
   void triggerUI();
 
+  void invalidate();
+
   jni::global_ref<AndroidUIScheduler::javaobject> javaPart_;
   std::shared_ptr<UIScheduler> uiScheduler_;
 

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

@@ -18,8 +18,15 @@ public class AndroidUIScheduler {
 
   private final Runnable mUIThreadRunnable =
       () -> {
-        if (mActive.get()) {
-          triggerUI();
+        /**
+         * This callback is called on the UI thread, but the module is invalidated on the JS thread.
+         * Therefore we must synchronize for reloads. Without synchronization the cpp part gets torn
+         * down while the UI thread is still executing it, leading to crashes.
+         */
+        synchronized (mActive) {
+          if (mActive.get()) {
+            triggerUI();
+          }
         }
       };
 
@@ -32,6 +39,8 @@ public AndroidUIScheduler(ReactApplicationContext context) {
 
   public native void triggerUI();
 
+  public native void invalidate();
+
   @DoNotStrip
   private void scheduleTriggerOnUI() {
     UiThreadUtil.runOnUiThread(
@@ -43,6 +52,9 @@ public void runGuarded() {
   }
 
   public void deactivate() {
-    mActive.set(false);
+    synchronized (mActive) {
+      mActive.set(false);
+      invalidate();
+    }
   }
 }