Fix AnimatedPropsRegistry surface-stop resurrection race (#57376)

Summary:

AnimatedPropsRegistry::update() runs on the UI thread every animation frame and
created a surface's entry via operator[]. clearOnSurfaceStop() (run on the JS
thread when a surface stops) erases that entry, but an in-flight animation frame
landing after the stop re-created it via operator[] -- and since the surface is
gone, nothing ever cleans it up again. The resurrected entry leaks its
PropsSnapshot and ShadowNodeFamily for the lifetime of the registry.

A surface's entry is legitimately created by getMap(), which
AnimationBackendCommitHook calls on every React commit. stopSurface drains
in-flight commits before unregistering the ShadowTree
(ShadowTreeRegistry::remove takes the registry's unique lock, which excludes the
shared-locked commit visits), so getMap() can never run for a stopped surface.
That leaves update()'s operator[] as the only thing that can resurrect one.

Fix: update() now only refines surfaces that already exist (find instead of
operator[]) and never creates an entry; getMap() remains the sole creator. A
stopped surface can no longer be resurrected, and there is no extra bookkeeping
that could grow over time.

Changelog: [General][Fixed] - Fix a surface-stop race in the C++ Animated shared backend that could permanently leak per-surface animated state

Differential Revision: D109156094

Author: Bartlomiej Bloniarz

packages/react-native/ReactCommon/react/renderer/animated/NativeAnimatedNodesManagerProvider.cpp

@@ -119,7 +119,7 @@ NativeAnimatedNodesManagerProvider::getOrCreate(
         });
 
     // Register on surfaces started in the future
-    uiManager->setOnSurfaceStartCallback(
+    uiManager->addOnSurfaceStartCallback(
         [animatedMountingOverrideDelegate =
              std::weak_ptr<const AnimatedMountingOverrideDelegate>(
                  animatedMountingOverrideDelegate_)](

packages/react-native/ReactCommon/react/renderer/animationbackend/AnimatedPropsRegistry.cpp

@@ -15,7 +15,11 @@ void AnimatedPropsRegistry::update(
     const std::unordered_map<SurfaceId, SurfaceUpdates>& surfaceUpdates) {
   auto lock = std::lock_guard(mutex_);
   for (const auto& [surfaceId, updates] : surfaceUpdates) {
-    auto& surfaceContext = surfaceContexts_[surfaceId];
+    auto contextIt = surfaceContexts_.find(surfaceId);
+    if (contextIt == surfaceContexts_.end()) {
+      continue;
+    }
+    auto& surfaceContext = contextIt->second;
     auto& pendingMap = surfaceContext.pendingMap;
     auto& pendingFamilies = surfaceContext.pendingFamilies;
 
@@ -55,6 +59,11 @@ void AnimatedPropsRegistry::update(
   }
 }
 
+void AnimatedPropsRegistry::initializeSurface(SurfaceId surfaceId) {
+  auto lock = std::lock_guard(mutex_);
+  surfaceContexts_.try_emplace(surfaceId);
+}
+
 std::pair<
     std::unordered_set<std::shared_ptr<const ShadowNodeFamily>>&,
     SnapshotMap&>

packages/react-native/ReactCommon/react/renderer/animationbackend/AnimatedPropsRegistry.h

@@ -38,6 +38,7 @@ using SnapshotMap = std::unordered_map<Tag, std::unique_ptr<PropsSnapshot>>;
 class AnimatedPropsRegistry {
  public:
   void update(const std::unordered_map<SurfaceId, SurfaceUpdates> &surfaceUpdates);
+  void initializeSurface(SurfaceId surfaceId);
   void clear(SurfaceId surfaceId);
   void clearOnSurfaceStop(SurfaceId surfaceId);
   std::pair<std::unordered_set<std::shared_ptr<const ShadowNodeFamily>> &, SnapshotMap &> getMap(SurfaceId surfaceId);

packages/react-native/ReactCommon/react/renderer/animationbackend/AnimationBackend.cpp

@@ -56,6 +56,24 @@ AnimationBackend::AnimationBackend(
       commitHook_(*uiManager, animatedPropsRegistry_),
       uiManager_(std::move(uiManager)) {
   react_native_assert(uiManager_.expired() == false);
+
+  auto weakAnimatedPropsRegistry =
+      std::weak_ptr<AnimatedPropsRegistry>(animatedPropsRegistry_);
+  auto initializeSurfaceContext =
+      [weakAnimatedPropsRegistry](const ShadowTree& shadowTree) {
+        if (auto animatedPropsRegistry = weakAnimatedPropsRegistry.lock()) {
+          animatedPropsRegistry->initializeSurface(shadowTree.getSurfaceId());
+        }
+      };
+
+  if (auto lockedUIManager = uiManager_.lock()) {
+    lockedUIManager->addOnSurfaceStartCallback(initializeSurfaceContext);
+    lockedUIManager->getShadowTreeRegistry().enumerate(
+        [&initializeSurfaceContext](
+            const ShadowTree& shadowTree, bool& /*stop*/) {
+          initializeSurfaceContext(shadowTree);
+        });
+  }
 }
 
 void AnimationBackend::unpackMutations(

packages/react-native/ReactCommon/react/renderer/scheduler/Scheduler.cpp

@@ -22,6 +22,7 @@
 #include <react/renderer/runtimescheduler/RuntimeScheduler.h>
 #include <react/renderer/uimanager/UIManager.h>
 #include <react/renderer/uimanager/UIManagerBinding.h>
+#include <mutex>
 
 namespace facebook::react {
 
@@ -459,9 +460,13 @@ void Scheduler::uiManagerDidPromoteReactRevision(const ShadowTree& shadowTree) {
 }
 
 void Scheduler::uiManagerDidStartSurface(const ShadowTree& shadowTree) {
-  std::shared_lock lock(onSurfaceStartCallbackMutex_);
-  if (onSurfaceStartCallback_) {
-    onSurfaceStartCallback_(shadowTree);
+  std::vector<OnSurfaceStartCallback> callbacks;
+  {
+    std::shared_lock lock(onSurfaceStartCallbackMutex_);
+    callbacks = onSurfaceStartCallbacks_;
+  }
+  for (const auto& callback : callbacks) {
+    callback(shadowTree);
   }
 }
 
@@ -491,10 +496,10 @@ void Scheduler::removeEventListener(
   }
 }
 
-void Scheduler::uiManagerShouldSetOnSurfaceStartCallback(
+void Scheduler::uiManagerShouldAddOnSurfaceStartCallback(
     OnSurfaceStartCallback&& callback) {
-  std::shared_lock lock(onSurfaceStartCallbackMutex_);
-  onSurfaceStartCallback_ = std::move(callback);
+  std::unique_lock lock(onSurfaceStartCallbackMutex_);
+  onSurfaceStartCallbacks_.push_back(std::move(callback));
 }
 
 } // namespace facebook::react

packages/react-native/ReactCommon/react/renderer/scheduler/Scheduler.h

@@ -9,6 +9,7 @@
 
 #include <atomic>
 #include <memory>
+#include <vector>
 
 #include <ReactCommon/RuntimeExecutor.h>
 #include <react/performance/cdpmetrics/CdpMetricsReporter.h>
@@ -117,7 +118,7 @@ class Scheduler final : public UIManagerDelegate {
   void removeEventListener(const std::shared_ptr<const EventListener> &listener);
 
 #pragma mark - Surface start callback
-  void uiManagerShouldSetOnSurfaceStartCallback(OnSurfaceStartCallback &&callback) override;
+  void uiManagerShouldAddOnSurfaceStartCallback(OnSurfaceStartCallback &&callback) override;
 
  private:
   friend class SurfaceHandler;
@@ -159,7 +160,7 @@ class Scheduler final : public UIManagerDelegate {
   std::shared_ptr<ViewTransitionModule> viewTransitionModule_;
 
   mutable std::shared_mutex onSurfaceStartCallbackMutex_;
-  OnSurfaceStartCallback onSurfaceStartCallback_;
+  std::vector<OnSurfaceStartCallback> onSurfaceStartCallbacks_;
 };
 
 } // namespace facebook::react

packages/react-native/ReactCommon/react/renderer/uimanager/UIManager.cpp

@@ -765,10 +765,10 @@ void UIManager::removeEventListener(
   }
 }
 
-void UIManager::setOnSurfaceStartCallback(
+void UIManager::addOnSurfaceStartCallback(
     UIManagerDelegate::OnSurfaceStartCallback&& callback) {
   if (delegate_ != nullptr) {
-    delegate_->uiManagerShouldSetOnSurfaceStartCallback(std::move(callback));
+    delegate_->uiManagerShouldAddOnSurfaceStartCallback(std::move(callback));
   }
 }
 

packages/react-native/ReactCommon/react/renderer/uimanager/UIManager.h

@@ -223,8 +223,8 @@ class UIManager final : public ShadowTreeDelegate {
 
   void removeEventListener(const std::shared_ptr<const EventListener> &listener);
 
-#pragma mark - Set on surface start callback
-  void setOnSurfaceStartCallback(UIManagerDelegate::OnSurfaceStartCallback &&callback);
+#pragma mark - Add on surface start callback
+  void addOnSurfaceStartCallback(UIManagerDelegate::OnSurfaceStartCallback &&callback);
 
  private:
   friend class UIManagerBinding;

packages/react-native/ReactCommon/react/renderer/uimanager/UIManagerDelegate.h

@@ -95,7 +95,7 @@ class UIManagerDelegate {
   virtual void uiManagerDidPromoteReactRevision(const ShadowTree &shadowTree) = 0;
 
   using OnSurfaceStartCallback = std::function<void(const ShadowTree &shadowTree)>;
-  virtual void uiManagerShouldSetOnSurfaceStartCallback(OnSurfaceStartCallback &&callback) = 0;
+  virtual void uiManagerShouldAddOnSurfaceStartCallback(OnSurfaceStartCallback &&callback) = 0;
 
   // View transition bitmap snapshot capture and application.
   virtual void uiManagerDidCaptureViewSnapshot(Tag tag, SurfaceId surfaceId) = 0;

packages/react-native/ReactCommon/react/renderer/viewtransition/ViewTransitionModule.cpp

@@ -59,7 +59,7 @@ void ViewTransitionModule::initialize(
         });
 
     // Register on surfaces started in the future
-    uiManager_->setOnSurfaceStartCallback(
+    uiManager_->addOnSurfaceStartCallback(
         [weakThis](const ShadowTree& shadowTree) {
           shadowTree.getMountingCoordinator()->setMountingOverrideDelegate(
               weakThis);