Implement re-entrancy and unlocked scenarios for LazyShadowTreeRevisionConsistencyManager (#44620)

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

We added a log message when trying to lock revisions in `LazyShadowTreeRevisionConsistencyManager` when they were already locked, and we've seen that being logged in existing experiments, which could indicate we're doing re-entrance from the JS runtime.

This protects against that case migrating the boolean flag to an integer.

Changelog: [internal]

Reviewed By: NickGerleman

Differential Revision: D57509193

fbshipit-source-id: 1712aa84d665c9dfe50630818e7f56de7d7e145c

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

@@ -19,7 +19,13 @@ void LazyShadowTreeRevisionConsistencyManager::updateCurrentRevision(
     SurfaceId surfaceId,
     RootShadowNode::Shared rootShadowNode) {
   std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
-  capturedRootShadowNodesForConsistency_[surfaceId] = std::move(rootShadowNode);
+
+  // We don't need to store the revision if we haven't locked.
+  // We can resolve lazily when requested.
+  if (lockCount > 0) {
+    capturedRootShadowNodesForConsistency_[surfaceId] =
+        std::move(rootShadowNode);
+  }
 }
 
 #pragma mark - ShadowTreeRevisionProvider
@@ -29,10 +35,11 @@ LazyShadowTreeRevisionConsistencyManager::getCurrentRevision(
     SurfaceId surfaceId) {
   {
     std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
-
-    auto it = capturedRootShadowNodesForConsistency_.find(surfaceId);
-    if (it != capturedRootShadowNodesForConsistency_.end()) {
-      return it->second;
+    if (lockCount > 0) {
+      auto it = capturedRootShadowNodesForConsistency_.find(surfaceId);
+      if (it != capturedRootShadowNodesForConsistency_.end()) {
+        return it->second;
+      }
     }
   }
 
@@ -47,7 +54,9 @@ LazyShadowTreeRevisionConsistencyManager::getCurrentRevision(
 
   {
     std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
-    capturedRootShadowNodesForConsistency_[surfaceId] = rootShadowNode;
+    if (lockCount > 0) {
+      capturedRootShadowNodesForConsistency_[surfaceId] = rootShadowNode;
+    }
   }
 
   return rootShadowNode;
@@ -56,29 +65,26 @@ LazyShadowTreeRevisionConsistencyManager::getCurrentRevision(
 #pragma mark - ConsistentShadowTreeRevisionProvider
 
 void LazyShadowTreeRevisionConsistencyManager::lockRevisions() {
-  if (isLocked_) {
-    LOG(WARNING)
-        << "LazyShadowTreeRevisionConsistencyManager::lockRevisions() called without unlocking a previous lock";
-    return;
-  }
+  std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
 
   // We actually capture the state lazily the first time we access it, so we
   // don't need to do anything here.
-  isLocked_ = true;
+  lockCount++;
 }
 
 void LazyShadowTreeRevisionConsistencyManager::unlockRevisions() {
-  if (!isLocked_) {
+  std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
+
+  if (lockCount == 0) {
     LOG(WARNING)
         << "LazyShadowTreeRevisionConsistencyManager::unlockRevisions() called without a previous lock";
-    // We don't return here because we want to do the cleanup anyway
-    // to free up resources.
+  } else {
+    lockCount--;
   }
 
-  isLocked_ = false;
-
-  std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
-  capturedRootShadowNodesForConsistency_.clear();
+  if (lockCount == 0) {
+    capturedRootShadowNodesForConsistency_.clear();
+  }
 }
 
 } // namespace facebook::react

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

@@ -11,6 +11,7 @@
 #include <react/renderer/consistency/ShadowTreeRevisionConsistencyManager.h>
 #include <react/renderer/mounting/ShadowTreeRegistry.h>
 #include <react/renderer/uimanager/consistency/ShadowTreeRevisionProvider.h>
+#include <cstdint>
 #include <memory>
 #include <shared_mutex>
 
@@ -47,7 +48,7 @@ class LazyShadowTreeRevisionConsistencyManager
   std::unordered_map<SurfaceId, RootShadowNode::Shared>
       capturedRootShadowNodesForConsistency_;
   ShadowTreeRegistry& shadowTreeRegistry_;
-  bool isLocked_{false};
+  uint_fast32_t lockCount{0};
 };
 
 } // namespace facebook::react

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

@@ -93,6 +93,27 @@ TEST_F(LazyShadowTreeRevisionConsistencyManagerTest, testLockedOnNoRevision) {
   consistencyManager_.unlockRevisions();
 }
 
+TEST_F(LazyShadowTreeRevisionConsistencyManagerTest, testNotLocked) {
+  EXPECT_EQ(consistencyManager_.getCurrentRevision(0), nullptr);
+
+  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;
+            },
+            {});
+      });
+
+  EXPECT_EQ(consistencyManager_.getCurrentRevision(0), newRootShadowNode);
+}
+
 TEST_F(
     LazyShadowTreeRevisionConsistencyManagerTest,
     testLockedOnNoRevisionWithUpdate) {
@@ -330,4 +351,45 @@ TEST_F(LazyShadowTreeRevisionConsistencyManagerTest, testUpdateToUnmounted) {
   consistencyManager_.unlockRevisions();
 }
 
+TEST_F(LazyShadowTreeRevisionConsistencyManagerTest, testReentrance) {
+  consistencyManager_.lockRevisions();
+
+  EXPECT_EQ(consistencyManager_.getCurrentRevision(0), nullptr);
+
+  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;
+            },
+            {});
+      });
+
+  EXPECT_EQ(consistencyManager_.getCurrentRevision(0), nullptr);
+
+  // Re-entrance
+  consistencyManager_.lockRevisions();
+
+  EXPECT_EQ(consistencyManager_.getCurrentRevision(0), nullptr);
+
+  // Exit second lock
+  consistencyManager_.unlockRevisions();
+
+  EXPECT_EQ(consistencyManager_.getCurrentRevision(0), nullptr);
+
+  // Exit first lock
+
+  consistencyManager_.unlockRevisions();
+
+  // Updated!
+  EXPECT_EQ(
+      consistencyManager_.getCurrentRevision(0).get(), newRootShadowNode.get());
+}
+
 } // namespace facebook::react