fix: Memory mismanagement with UI scheduled callbacks (#6900)

The pattern we used with scheduling callbacks on the UI thread -
capturing `this` as a reference - is prone to memory issues. Given such
code:

```c++
void WorkletsModuleProxy::scheduleOnUI(
    jsi::Runtime &rt,
    const jsi::Value &worklet) {
  auto shareableWorklet = extractShareableOrThrow<ShareableWorklet>(
      rt, worklet, "[Worklets] Only worklets can be scheduled to run on UI.");
  uiScheduler_->scheduleOnUI(
      [this, shareableWorklet] {
        this->uiWorkletRuntime_->runGuarded(shareableWorklet);
      });
}
```

We are likely to run into accessing invalidated memory during a reload.
This is due to fact that `WorkletsModuleProxy` is managed by some object
held by the instance of React Native. Let's look at the following
scenario.

1. `WorkletsModuleProxy` is created on the JS thread and held by the
`WorkletsModule` Native Module.
2. `WorkletsModuleProxy::scheduleOnUI` is invoked on the JS thread. The
callback is scheduled to be executed on the UI thread.
3. Application's reload gets triggered. A tear down of React Native is
starting on the JS thread.
4. `WorkletsModule` gets destroyed. Therefore, `WorkletsModuleProxy` is
released and also destroyed.
5. The callback is finally executed on the UI thread by the scheduler.
However, `this` has been invalidated. The App crashes.

Keep in mind that this isn't exclusive to thread jumps exclusively.
Calling `scheduleOnUI` on the UI thread could still result in the
callback executing after the memory has been invalidated.

`WorkletsModuleProxy` is only an example here, the problem could
manifest in all the places where we pass lambdas that capture `this` by
reference.

To fix this I refactored the code so everytime we pass `this` to a
scheduled callback, it would be done via a _weak pointer_ which would
lock the object and prevent it from being destroyed while the callback
is being executed on the UI thread.

Perhaps some bits of code don't need this safety measure due to a
heuristic existing that guarantees that respective memory won't be
invalidated before the callback gets executed. However, I found it
extremely challenging and unreliable to come up with these heuristics,
as they could possibly break at any future change of the code.

- ReanimatedCommitHook
- LayoutAnimationProxy
- ReanimatedModuleProxy
- WorkletsModuleProxy
- WorkletRuntime
- NativeProxy

Reloading the app no longer causes a crash on a scheduled UI callback.

Author: tjzel

packages/react-native-reanimated/Common/cpp/reanimated/Fabric/ReanimatedCommitHook.cpp

@@ -34,14 +34,16 @@ void ReanimatedCommitHook::maybeInitializeLayoutAnimations(
     // when a new surfaceId is observed we call setMountingOverrideDelegate
     // for all yet unseen surfaces
     uiManager_->getShadowTreeRegistry().enumerate(
-        [this](const ShadowTree &shadowTree, bool &stop) {
-          if (shadowTree.getSurfaceId() <= currentMaxSurfaceId_) {
+        [strongThis = shared_from_this()](
+            const ShadowTree &shadowTree, bool &stop) {
+          // Executed synchronously.
+          if (shadowTree.getSurfaceId() <= strongThis->currentMaxSurfaceId_) {
             // the set function actually adds our delegate to a list, so we
             // shouldn't invoke it twice for the same surface
             return;
           }
           shadowTree.getMountingCoordinator()->setMountingOverrideDelegate(
-              layoutAnimationsProxy_);
+              strongThis->layoutAnimationsProxy_);
         });
     currentMaxSurfaceId_ = surfaceId;
   }

packages/react-native-reanimated/Common/cpp/reanimated/Fabric/ReanimatedCommitHook.h

@@ -12,7 +12,9 @@ using namespace facebook::react;
 
 namespace reanimated {
 
-class ReanimatedCommitHook : public UIManagerCommitHook {
+class ReanimatedCommitHook
+    : public UIManagerCommitHook,
+      public std::enable_shared_from_this<ReanimatedCommitHook> {
  public:
   ReanimatedCommitHook(
       const std::shared_ptr<PropsRegistry> &propsRegistry,

packages/react-native-reanimated/Common/cpp/reanimated/LayoutAnimations/LayoutAnimationsProxy.cpp

@@ -650,19 +650,25 @@ void LayoutAnimationsProxy::startEnteringAnimation(
       static_cast<const ViewProps &>(*mutation.newChildShadowView.props);
   auto opacity = viewProps.opacity;
 
-  uiScheduler_->scheduleOnUI([finalView,
+  uiScheduler_->scheduleOnUI([weakThis = weak_from_this(),
+                              finalView,
                               current,
 #if REACT_NATIVE_MINOR_VERSION < 78
                               parent,
-#endif // REACT_NATIVE_MINOR_VERSION < 78
+#endif // RE
                               mutation,
                               opacity,
-                              this,
                               tag]() {
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
+
     Rect window{};
     {
+      auto &mutex = strongThis->mutex;
       auto lock = std::unique_lock<std::recursive_mutex>(mutex);
-      layoutAnimations_.insert_or_assign(
+      strongThis->layoutAnimations_.insert_or_assign(
           tag, LayoutAnimation {
             finalView, current,
 #if REACT_NATIVE_MINOR_VERSION >= 78
@@ -672,21 +678,23 @@ void LayoutAnimationsProxy::startEnteringAnimation(
 #endif // REACT_NATIVE_MINOR_VERSION >= 78
                 opacity
           });
-      window = surfaceManager.getWindow(mutation.newChildShadowView.surfaceId);
+      window = strongThis->surfaceManager.getWindow(
+          mutation.newChildShadowView.surfaceId);
     }
 
     Snapshot values(mutation.newChildShadowView, window);
-    jsi::Object yogaValues(uiRuntime_);
-    yogaValues.setProperty(uiRuntime_, "targetOriginX", values.x);
-    yogaValues.setProperty(uiRuntime_, "targetGlobalOriginX", values.x);
-    yogaValues.setProperty(uiRuntime_, "targetOriginY", values.y);
-    yogaValues.setProperty(uiRuntime_, "targetGlobalOriginY", values.y);
-    yogaValues.setProperty(uiRuntime_, "targetWidth", values.width);
-    yogaValues.setProperty(uiRuntime_, "targetHeight", values.height);
-    yogaValues.setProperty(uiRuntime_, "windowWidth", values.windowWidth);
-    yogaValues.setProperty(uiRuntime_, "windowHeight", values.windowHeight);
-    layoutAnimationsManager_->startLayoutAnimation(
-        uiRuntime_, tag, LayoutAnimationType::ENTERING, yogaValues);
+    auto &uiRuntime = strongThis->uiRuntime_;
+    jsi::Object yogaValues(uiRuntime);
+    yogaValues.setProperty(uiRuntime, "targetOriginX", values.x);
+    yogaValues.setProperty(uiRuntime, "targetGlobalOriginX", values.x);
+    yogaValues.setProperty(uiRuntime, "targetOriginY", values.y);
+    yogaValues.setProperty(uiRuntime, "targetGlobalOriginY", values.y);
+    yogaValues.setProperty(uiRuntime, "targetWidth", values.width);
+    yogaValues.setProperty(uiRuntime, "targetHeight", values.height);
+    yogaValues.setProperty(uiRuntime, "windowWidth", values.windowWidth);
+    yogaValues.setProperty(uiRuntime, "windowHeight", values.windowHeight);
+    strongThis->layoutAnimationsManager_->startLayoutAnimation(
+        uiRuntime, tag, LayoutAnimationType::ENTERING, yogaValues);
   });
 }
 
@@ -698,30 +706,38 @@ void LayoutAnimationsProxy::startExitingAnimation(
 #endif
   auto surfaceId = mutation.oldChildShadowView.surfaceId;
 
-  uiScheduler_->scheduleOnUI([this, tag, mutation, surfaceId]() {
-    auto oldView = mutation.oldChildShadowView;
-    Rect window{};
-    {
-      auto lock = std::unique_lock<std::recursive_mutex>(mutex);
-      createLayoutAnimation(mutation, oldView, surfaceId, tag);
-      window = surfaceManager.getWindow(surfaceId);
-    }
+  uiScheduler_->scheduleOnUI(
+      [weakThis = weak_from_this(), tag, mutation, surfaceId]() {
+        auto strongThis = weakThis.lock();
+        if (!strongThis) {
+          return;
+        }
 
-    Snapshot values(oldView, window);
-
-    jsi::Object yogaValues(uiRuntime_);
-    yogaValues.setProperty(uiRuntime_, "currentOriginX", values.x);
-    yogaValues.setProperty(uiRuntime_, "currentGlobalOriginX", values.x);
-    yogaValues.setProperty(uiRuntime_, "currentOriginY", values.y);
-    yogaValues.setProperty(uiRuntime_, "currentGlobalOriginY", values.y);
-    yogaValues.setProperty(uiRuntime_, "currentWidth", values.width);
-    yogaValues.setProperty(uiRuntime_, "currentHeight", values.height);
-    yogaValues.setProperty(uiRuntime_, "windowWidth", values.windowWidth);
-    yogaValues.setProperty(uiRuntime_, "windowHeight", values.windowHeight);
-    layoutAnimationsManager_->startLayoutAnimation(
-        uiRuntime_, tag, LayoutAnimationType::EXITING, yogaValues);
-    layoutAnimationsManager_->clearLayoutAnimationConfig(tag);
-  });
+        auto oldView = mutation.oldChildShadowView;
+        Rect window{};
+        {
+          auto &mutex = strongThis->mutex;
+          auto lock = std::unique_lock<std::recursive_mutex>(mutex);
+          strongThis->createLayoutAnimation(mutation, oldView, surfaceId, tag);
+          window = strongThis->surfaceManager.getWindow(surfaceId);
+        }
+
+        Snapshot values(oldView, window);
+
+        auto &uiRuntime = strongThis->uiRuntime_;
+        jsi::Object yogaValues(uiRuntime);
+        yogaValues.setProperty(uiRuntime, "currentOriginX", values.x);
+        yogaValues.setProperty(uiRuntime, "currentGlobalOriginX", values.x);
+        yogaValues.setProperty(uiRuntime, "currentOriginY", values.y);
+        yogaValues.setProperty(uiRuntime, "currentGlobalOriginY", values.y);
+        yogaValues.setProperty(uiRuntime, "currentWidth", values.width);
+        yogaValues.setProperty(uiRuntime, "currentHeight", values.height);
+        yogaValues.setProperty(uiRuntime, "windowWidth", values.windowWidth);
+        yogaValues.setProperty(uiRuntime, "windowHeight", values.windowHeight);
+        strongThis->layoutAnimationsManager_->startLayoutAnimation(
+            uiRuntime, tag, LayoutAnimationType::EXITING, yogaValues);
+        strongThis->layoutAnimationsManager_->clearLayoutAnimationConfig(tag);
+      });
 }
 
 void LayoutAnimationsProxy::startLayoutAnimation(
@@ -732,36 +748,46 @@ void LayoutAnimationsProxy::startLayoutAnimation(
 #endif
   auto surfaceId = mutation.oldChildShadowView.surfaceId;
 
-  uiScheduler_->scheduleOnUI([this, mutation, surfaceId, tag]() {
+  uiScheduler_->scheduleOnUI([weakThis = weak_from_this(),
+                              mutation,
+                              surfaceId,
+                              tag]() {
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
+
     auto oldView = mutation.oldChildShadowView;
     Rect window{};
     {
+      auto &mutex = strongThis->mutex;
       auto lock = std::unique_lock<std::recursive_mutex>(mutex);
-      createLayoutAnimation(mutation, oldView, surfaceId, tag);
-      window = surfaceManager.getWindow(surfaceId);
+      strongThis->createLayoutAnimation(mutation, oldView, surfaceId, tag);
+      window = strongThis->surfaceManager.getWindow(surfaceId);
     }
 
     Snapshot currentValues(oldView, window);
     Snapshot targetValues(mutation.newChildShadowView, window);
 
-    jsi::Object yogaValues(uiRuntime_);
-    yogaValues.setProperty(uiRuntime_, "currentOriginX", currentValues.x);
-    yogaValues.setProperty(uiRuntime_, "currentGlobalOriginX", currentValues.x);
-    yogaValues.setProperty(uiRuntime_, "currentOriginY", currentValues.y);
-    yogaValues.setProperty(uiRuntime_, "currentGlobalOriginY", currentValues.y);
-    yogaValues.setProperty(uiRuntime_, "currentWidth", currentValues.width);
-    yogaValues.setProperty(uiRuntime_, "currentHeight", currentValues.height);
-    yogaValues.setProperty(uiRuntime_, "targetOriginX", targetValues.x);
-    yogaValues.setProperty(uiRuntime_, "targetGlobalOriginX", targetValues.x);
-    yogaValues.setProperty(uiRuntime_, "targetOriginY", targetValues.y);
-    yogaValues.setProperty(uiRuntime_, "targetGlobalOriginY", targetValues.y);
-    yogaValues.setProperty(uiRuntime_, "targetWidth", targetValues.width);
-    yogaValues.setProperty(uiRuntime_, "targetHeight", targetValues.height);
-    yogaValues.setProperty(uiRuntime_, "windowWidth", targetValues.windowWidth);
+    auto &uiRuntime = strongThis->uiRuntime_;
+    jsi::Object yogaValues(uiRuntime);
+    yogaValues.setProperty(uiRuntime, "currentOriginX", currentValues.x);
+    yogaValues.setProperty(uiRuntime, "currentGlobalOriginX", currentValues.x);
+    yogaValues.setProperty(uiRuntime, "currentOriginY", currentValues.y);
+    yogaValues.setProperty(uiRuntime, "currentGlobalOriginY", currentValues.y);
+    yogaValues.setProperty(uiRuntime, "currentWidth", currentValues.width);
+    yogaValues.setProperty(uiRuntime, "currentHeight", currentValues.height);
+    yogaValues.setProperty(uiRuntime, "targetOriginX", targetValues.x);
+    yogaValues.setProperty(uiRuntime, "targetGlobalOriginX", targetValues.x);
+    yogaValues.setProperty(uiRuntime, "targetOriginY", targetValues.y);
+    yogaValues.setProperty(uiRuntime, "targetGlobalOriginY", targetValues.y);
+    yogaValues.setProperty(uiRuntime, "targetWidth", targetValues.width);
+    yogaValues.setProperty(uiRuntime, "targetHeight", targetValues.height);
+    yogaValues.setProperty(uiRuntime, "windowWidth", targetValues.windowWidth);
     yogaValues.setProperty(
-        uiRuntime_, "windowHeight", targetValues.windowHeight);
-    layoutAnimationsManager_->startLayoutAnimation(
-        uiRuntime_, tag, LayoutAnimationType::LAYOUT, yogaValues);
+        uiRuntime, "windowHeight", targetValues.windowHeight);
+    strongThis->layoutAnimationsManager_->startLayoutAnimation(
+        uiRuntime, tag, LayoutAnimationType::LAYOUT, yogaValues);
   });
 }
 
@@ -777,8 +803,14 @@ void LayoutAnimationsProxy::maybeCancelAnimation(const int tag) const {
     return;
   }
   layoutAnimations_.erase(tag);
-  uiScheduler_->scheduleOnUI([this, tag]() {
-    layoutAnimationsManager_->cancelLayoutAnimation(uiRuntime_, tag);
+  uiScheduler_->scheduleOnUI([weakThis = weak_from_this(), tag]() {
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
+
+    auto &uiRuntime = strongThis->uiRuntime_;
+    strongThis->layoutAnimationsManager_->cancelLayoutAnimation(uiRuntime, tag);
   });
 }
 

packages/react-native-reanimated/Common/cpp/reanimated/LayoutAnimations/LayoutAnimationsProxy.h

@@ -35,7 +35,9 @@ struct LayoutAnimation {
   LayoutAnimation &operator=(const LayoutAnimation &other) = default;
 };
 
-struct LayoutAnimationsProxy : public MountingOverrideDelegate {
+struct LayoutAnimationsProxy
+    : public MountingOverrideDelegate,
+      public std::enable_shared_from_this<LayoutAnimationsProxy> {
   mutable std::unordered_map<Tag, std::shared_ptr<Node>> nodeForTag_;
   mutable std::unordered_map<Tag, LayoutAnimation> layoutAnimations_;
   mutable std::recursive_mutex mutex;