Fix AnimatedPropsRegistry surface-stop resurrection race

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/animationbackend/AnimatedPropsRegistry.cpp

@@ -15,7 +15,12 @@ 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];
+    // Do not recreate state for stopped surfaces from late animation frames.
+    auto contextIt = surfaceContexts_.find(surfaceId);
+    if (contextIt == surfaceContexts_.end()) {
+      continue;
+    }
+    auto& surfaceContext = contextIt->second;
     auto& pendingMap = surfaceContext.pendingMap;
     auto& pendingFamilies = surfaceContext.pendingFamilies;
 
@@ -55,6 +60,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,25 @@ 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->getShadowTreeRegistry().enumerate(
+        [&initializeSurfaceContext](
+            const ShadowTree& shadowTree, bool& /*stop*/) {
+          initializeSurfaceContext(shadowTree);
+        });
+    lockedUIManager->setOnSurfaceStartCallback(
+        std::move(initializeSurfaceContext));
+  }
 }
 
 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 {
 
@@ -493,8 +494,21 @@ void Scheduler::removeEventListener(
 
 void Scheduler::uiManagerShouldSetOnSurfaceStartCallback(
     OnSurfaceStartCallback&& callback) {
-  std::shared_lock lock(onSurfaceStartCallbackMutex_);
-  onSurfaceStartCallback_ = std::move(callback);
+  std::unique_lock lock(onSurfaceStartCallbackMutex_);
+  if (!onSurfaceStartCallback_) {
+    onSurfaceStartCallback_ = std::move(callback);
+    return;
+  }
+
+  auto previousCallback = std::move(onSurfaceStartCallback_);
+  onSurfaceStartCallback_ =
+      [previousCallback = std::move(previousCallback),
+       nextCallback = std::move(callback)](const ShadowTree& shadowTree) {
+        previousCallback(shadowTree);
+        if (nextCallback) {
+          nextCallback(shadowTree);
+        }
+      };
 }
 
 } // namespace facebook::react

scripts/cxx-api/api-snapshots/ReactAndroidDebugCxx.api

@@ -1557,6 +1557,7 @@ class facebook::react::AnimatedPropsRegistry {
   public std::pair<std::unordered_set<std::shared_ptr<const facebook::react::ShadowNodeFamily>>&, facebook::react::SnapshotMap&> getMap(facebook::react::SurfaceId surfaceId);
   public void clear(facebook::react::SurfaceId surfaceId);
   public void clearOnSurfaceStop(facebook::react::SurfaceId surfaceId);
+  public void initializeSurface(facebook::react::SurfaceId surfaceId);
   public void update(const std::unordered_map<facebook::react::SurfaceId, facebook::react::SurfaceUpdates>& surfaceUpdates);
 }
 

scripts/cxx-api/api-snapshots/ReactAndroidNewarchCxx.api

@@ -1552,6 +1552,7 @@ class facebook::react::AnimatedPropsRegistry {
   public std::pair<std::unordered_set<std::shared_ptr<const facebook::react::ShadowNodeFamily>>&, facebook::react::SnapshotMap&> getMap(facebook::react::SurfaceId surfaceId);
   public void clear(facebook::react::SurfaceId surfaceId);
   public void clearOnSurfaceStop(facebook::react::SurfaceId surfaceId);
+  public void initializeSurface(facebook::react::SurfaceId surfaceId);
   public void update(const std::unordered_map<facebook::react::SurfaceId, facebook::react::SurfaceUpdates>& surfaceUpdates);
 }
 

scripts/cxx-api/api-snapshots/ReactAndroidReleaseCxx.api

@@ -1556,6 +1556,7 @@ class facebook::react::AnimatedPropsRegistry {
   public std::pair<std::unordered_set<std::shared_ptr<const facebook::react::ShadowNodeFamily>>&, facebook::react::SnapshotMap&> getMap(facebook::react::SurfaceId surfaceId);
   public void clear(facebook::react::SurfaceId surfaceId);
   public void clearOnSurfaceStop(facebook::react::SurfaceId surfaceId);
+  public void initializeSurface(facebook::react::SurfaceId surfaceId);
   public void update(const std::unordered_map<facebook::react::SurfaceId, facebook::react::SurfaceUpdates>& surfaceUpdates);
 }
 

scripts/cxx-api/api-snapshots/ReactAppleDebugCxx.api

@@ -4141,6 +4141,7 @@ class facebook::react::AnimatedPropsRegistry {
   public std::pair<std::unordered_set<std::shared_ptr<const facebook::react::ShadowNodeFamily>>&, facebook::react::SnapshotMap&> getMap(facebook::react::SurfaceId surfaceId);
   public void clear(facebook::react::SurfaceId surfaceId);
   public void clearOnSurfaceStop(facebook::react::SurfaceId surfaceId);
+  public void initializeSurface(facebook::react::SurfaceId surfaceId);
   public void update(const std::unordered_map<facebook::react::SurfaceId, facebook::react::SurfaceUpdates>& surfaceUpdates);
 }
 

scripts/cxx-api/api-snapshots/ReactAppleNewarchCxx.api

@@ -4125,6 +4125,7 @@ class facebook::react::AnimatedPropsRegistry {
   public std::pair<std::unordered_set<std::shared_ptr<const facebook::react::ShadowNodeFamily>>&, facebook::react::SnapshotMap&> getMap(facebook::react::SurfaceId surfaceId);
   public void clear(facebook::react::SurfaceId surfaceId);
   public void clearOnSurfaceStop(facebook::react::SurfaceId surfaceId);
+  public void initializeSurface(facebook::react::SurfaceId surfaceId);
   public void update(const std::unordered_map<facebook::react::SurfaceId, facebook::react::SurfaceUpdates>& surfaceUpdates);
 }
 

scripts/cxx-api/api-snapshots/ReactAppleReleaseCxx.api

@@ -4140,6 +4140,7 @@ class facebook::react::AnimatedPropsRegistry {
   public std::pair<std::unordered_set<std::shared_ptr<const facebook::react::ShadowNodeFamily>>&, facebook::react::SnapshotMap&> getMap(facebook::react::SurfaceId surfaceId);
   public void clear(facebook::react::SurfaceId surfaceId);
   public void clearOnSurfaceStop(facebook::react::SurfaceId surfaceId);
+  public void initializeSurface(facebook::react::SurfaceId surfaceId);
   public void update(const std::unordered_map<facebook::react::SurfaceId, facebook::react::SurfaceUpdates>& surfaceUpdates);
 }