Implement thread-safety for LazyShadowTreeRevisionConsistencyManager (#44619)

Summary:
Pull Request resolved: https://github.com/facebook/react-native/pull/44619

Some methods in `LazyShadowTreeRevisionConsistencyManager` can be called in parallel when using synchronous state updates (which is also behind a flag). This implements thread-safety to cover that case so we don't have issues when testing that variant in production.

Changelog: [internal]

Reviewed By: NickGerleman

Differential Revision: D57506540

fbshipit-source-id: 362e1df534bc8c87289882236cfe0d7ee261f507

packages/react-native/ReactCommon/react/renderer/uimanager/consistency/LazyShadowTreeRevisionConsistencyManager.cpp

@@ -18,6 +18,7 @@ LazyShadowTreeRevisionConsistencyManager::
 void LazyShadowTreeRevisionConsistencyManager::updateCurrentRevision(
     SurfaceId surfaceId,
     RootShadowNode::Shared rootShadowNode) {
+  std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
   capturedRootShadowNodesForConsistency_[surfaceId] = std::move(rootShadowNode);
 }
 
@@ -26,18 +27,28 @@ void LazyShadowTreeRevisionConsistencyManager::updateCurrentRevision(
 RootShadowNode::Shared
 LazyShadowTreeRevisionConsistencyManager::getCurrentRevision(
     SurfaceId surfaceId) {
-  auto it = capturedRootShadowNodesForConsistency_.find(surfaceId);
-  if (it != capturedRootShadowNodesForConsistency_.end()) {
-    return it->second;
+  {
+    std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
+
+    auto it = capturedRootShadowNodesForConsistency_.find(surfaceId);
+    if (it != capturedRootShadowNodesForConsistency_.end()) {
+      return it->second;
+    }
   }
 
+  // This method is only going to be called from JS, so we don't need to protect
+  // the access to the shadow tree registry as well.
+  // If this was multi-threaded, we would need to protect it to avoid capturing
+  // root shadow nodes concurrently.
   RootShadowNode::Shared rootShadowNode;
-
   shadowTreeRegistry_.visit(surfaceId, [&](const ShadowTree& shadowTree) {
     rootShadowNode = shadowTree.getCurrentRevision().rootShadowNode;
   });
 
-  capturedRootShadowNodesForConsistency_[surfaceId] = rootShadowNode;
+  {
+    std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
+    capturedRootShadowNodesForConsistency_[surfaceId] = rootShadowNode;
+  }
 
   return rootShadowNode;
 }
@@ -65,6 +76,8 @@ void LazyShadowTreeRevisionConsistencyManager::unlockRevisions() {
   }
 
   isLocked_ = false;
+
+  std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
   capturedRootShadowNodesForConsistency_.clear();
 }
 

packages/react-native/ReactCommon/react/renderer/uimanager/consistency/LazyShadowTreeRevisionConsistencyManager.h

@@ -12,6 +12,7 @@
 #include <react/renderer/mounting/ShadowTreeRegistry.h>
 #include <react/renderer/uimanager/consistency/ShadowTreeRevisionProvider.h>
 #include <memory>
+#include <shared_mutex>
 
 namespace facebook::react {
 
@@ -42,6 +43,7 @@ class LazyShadowTreeRevisionConsistencyManager
   void unlockRevisions() override;
 
  private:
+  std::mutex capturedRootShadowNodesForConsistencyMutex_;
   std::unordered_map<SurfaceId, RootShadowNode::Shared>
       capturedRootShadowNodesForConsistency_;
   ShadowTreeRegistry& shadowTreeRegistry_;

packages/react-native/ReactCommon/react/renderer/uimanager/consistency/tests/LazyShadowTreeRevisionConsistencyManagerTest.cpp

@@ -301,4 +301,33 @@ TEST_F(LazyShadowTreeRevisionConsistencyManagerTest, testLockAfterUnlock) {
   consistencyManager_.unlockRevisions();
 }
 
+TEST_F(LazyShadowTreeRevisionConsistencyManagerTest, testUpdateToUnmounted) {
+  shadowTreeRegistry_.add(createShadowTree(0));
+
+  auto element = Element<RootShadowNode>();
+  auto builder = simpleComponentBuilder();
+  auto newRootShadowNode = builder.build(element);
+
+  shadowTreeRegistry_.visit(
+      0, [newRootShadowNode](const ShadowTree& shadowTree) {
+        shadowTree.commit(
+            [&](const RootShadowNode& /*oldRootShadowNode*/) {
+              return newRootShadowNode;
+            },
+            {});
+      });
+
+  consistencyManager_.lockRevisions();
+
+  EXPECT_EQ(
+      consistencyManager_.getCurrentRevision(0).get(), newRootShadowNode.get());
+
+  consistencyManager_.updateCurrentRevision(0, nullptr);
+
+  // Updated
+  EXPECT_EQ(consistencyManager_.getCurrentRevision(0).get(), nullptr);
+
+  consistencyManager_.unlockRevisions();
+}
+
 } // namespace facebook::react