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

## Summary

### What
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.

### Fix
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.

### Affected code:

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

## Test plan

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

apps/fabric-example/ios/Podfile.lock

@@ -2254,7 +2254,7 @@ SPEC CHECKSUMS:
   RNScreens: d0854539b51a53e38b61bcc9fb402439a9c73b26
   RNSVG: 7e38044415125a1d108294377de261d2fe2c54c9
   SocketRocket: d4aabe649be1e368d1318fdf28a022d714d65748
-  Yoga: dd05f79b1943c96893294ee5040576aa6e8a8d82
+  Yoga: 0c8754b0ea9edb13b6ce6b60f0f69eb5f164f16a
 
 PODFILE CHECKSUM: 4a9e0af2552a3fcd2303b56ad75e373f8bae65b9
 

apps/macos-example/macos/Podfile.lock

@@ -1874,7 +1874,7 @@ SPEC CHECKSUMS:
   RNReanimated: c5abea1d20bfb7a2d3c7d3a1ce95a253bb9fb592
   RNSVG: 669ed128ab9005090c612a0d627dbecb6ab5c76f
   SocketRocket: 9ee265c4b5ae2382d18e4ee1d2dd2d7af0ff1ab5
-  Yoga: 209f62622a01344dbb9fa8d348610eaeb7df2cca
+  Yoga: 446e6f351a519539ff00a1159fe41e589aab1b94
 
 PODFILE CHECKSUM: 8d50cc2acc9f6a6b1a12bd9106b86385ad72266f
 

apps/tvos-example/ios/Podfile.lock

@@ -1888,7 +1888,7 @@ SPEC CHECKSUMS:
   ReactCommon: b927fd46115bd2acb146e24cf1a08f22abda8b3f
   RNReanimated: c5133742ac745dc11856e0b665e5acd65b4dac05
   SocketRocket: d4aabe649be1e368d1318fdf28a022d714d65748
-  Yoga: 2b0e5affb9ab46e4ebad33530df829c153c323d8
+  Yoga: 651e5fd560c7e408ab9d9ca44b8de1b622d7f0cc
 
 PODFILE CHECKSUM: 79e1477a8eb76b717bdd7c1610f7f8e6772536a9
 

packages/react-native-reanimated/Common/cpp/reanimated/CSS/registry/CSSAnimationsRegistry.h

@@ -15,7 +15,9 @@
 
 namespace reanimated {
 
-class CSSAnimationsRegistry : public UpdatesRegistry {
+class CSSAnimationsRegistry
+    : public UpdatesRegistry,
+      std::enable_shared_from_this<CSSAnimationsRegistry> {
  public:
   using SettingsUpdates =
       std::vector<std::pair<unsigned, PartialCSSAnimationSettings>>;

packages/react-native-reanimated/Common/cpp/reanimated/CSS/registry/CSSTransitionsRegistry.cpp

@@ -118,11 +118,16 @@ void CSSTransitionsRegistry::scheduleOrActivateTransition(
 }
 
 PropsObserver CSSTransitionsRegistry::createPropsObserver(const Tag viewTag) {
-  return [this, viewTag](
+  return [weakThis = weak_from_this(), viewTag](
              jsi::Runtime &rt,
              const jsi::Value &oldProps,
              const jsi::Value &newProps) {
-    const auto &transition = registry_.at(viewTag);
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
+
+    const auto &transition = strongThis->registry_.at(viewTag);
     const auto changedProps = getChangedProps(
         rt,
         oldProps,
@@ -135,14 +140,14 @@ PropsObserver CSSTransitionsRegistry::createPropsObserver(const Tag viewTag) {
     }
 
     {
-      std::lock_guard<std::mutex> lock{mutex_};
+      std::lock_guard<std::mutex> lock{strongThis->mutex_};
 
       const auto &initialProps =
-          transition->run(rt, changedProps, getCurrentTimestamp_());
+          transition->run(rt, changedProps, strongThis->getCurrentTimestamp_());
       const auto &shadowNode = transition->getShadowNode();
 
-      setInUpdatesRegistry(rt, shadowNode, initialProps);
-      scheduleOrActivateTransition(transition);
+      strongThis->setInUpdatesRegistry(rt, shadowNode, initialProps);
+      strongThis->scheduleOrActivateTransition(transition);
     }
   };
 }

packages/react-native-reanimated/Common/cpp/reanimated/CSS/registry/CSSTransitionsRegistry.h

@@ -17,7 +17,9 @@
 
 namespace reanimated {
 
-class CSSTransitionsRegistry : public UpdatesRegistry {
+class CSSTransitionsRegistry
+    : public UpdatesRegistry,
+      public std::enable_shared_from_this<CSSTransitionsRegistry> {
  public:
   CSSTransitionsRegistry(
       const std::shared_ptr<StaticPropsRegistry> &staticPropsRegistry,

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<UIManager> &uiManager,

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

@@ -618,30 +618,42 @@ void LayoutAnimationsProxy::startEnteringAnimation(
       static_cast<const ViewProps &>(*mutation.newChildShadowView.props);
   auto opacity = viewProps.opacity;
 
-  uiScheduler_->scheduleOnUI(
-      [finalView, current, parent, mutation, opacity, this, tag]() {
-        Rect window{};
-        {
-          auto lock = std::unique_lock<std::recursive_mutex>(mutex);
-          layoutAnimations_.insert_or_assign(
-              tag, LayoutAnimation{finalView, current, parent, opacity});
-          window =
-              surfaceManager.getWindow(mutation.newChildShadowView.surfaceId);
-        }
+  uiScheduler_->scheduleOnUI([weakThis = weak_from_this(),
+                              finalView,
+                              current,
+                              parent,
+                              mutation,
+                              opacity,
+                              tag]() {
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
 
-        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);
-      });
+    Rect window{};
+    {
+      auto &mutex = strongThis->mutex;
+      auto lock = std::unique_lock<std::recursive_mutex>(mutex);
+      strongThis->layoutAnimations_.insert_or_assign(
+          tag, LayoutAnimation{finalView, current, parent, opacity});
+      window = strongThis->surfaceManager.getWindow(
+          mutation.newChildShadowView.surfaceId);
+    }
+
+    Snapshot values(mutation.newChildShadowView, window);
+    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);
+  });
 }
 
 void LayoutAnimationsProxy::startExitingAnimation(
@@ -652,30 +664,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(
@@ -686,36 +706,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);
   });
 }
 
@@ -731,8 +761,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

@@ -29,7 +29,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;