Implement re-entrance and unlocked scenarios for LazyShadowTreeRevisi…

…onConsistencyManager

Summary:
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.

Differential Revision: D57509193

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

@@ -18,20 +18,27 @@ LazyShadowTreeRevisionConsistencyManager::
 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) {
+    std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
+    capturedRootShadowNodesForConsistency_[surfaceId] =
+        std::move(rootShadowNode);
+  }
 }
 
 #pragma mark - ShadowTreeRevisionProvider
 
 RootShadowNode::Shared
 LazyShadowTreeRevisionConsistencyManager::getCurrentRevision(
     SurfaceId surfaceId) {
-  std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
+  if (lockCount > 0) {
+    std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
 
-  auto it = capturedRootShadowNodesForConsistency_.find(surfaceId);
-  if (it != capturedRootShadowNodesForConsistency_.end()) {
-    return it->second;
+    auto it = capturedRootShadowNodesForConsistency_.find(surfaceId);
+    if (it != capturedRootShadowNodesForConsistency_.end()) {
+      return it->second;
+    }
   }
 
   RootShadowNode::Shared rootShadowNode;
@@ -40,37 +47,33 @@ LazyShadowTreeRevisionConsistencyManager::getCurrentRevision(
     rootShadowNode = shadowTree.getCurrentRevision().rootShadowNode;
   });
 
-  capturedRootShadowNodesForConsistency_[surfaceId] = rootShadowNode;
+  if (lockCount > 0) {
+    capturedRootShadowNodesForConsistency_[surfaceId] = rootShadowNode;
+  }
 
   return rootShadowNode;
 }
 
 #pragma mark - ConsistentShadowTreeRevisionProvider
 
 void LazyShadowTreeRevisionConsistencyManager::lockRevisions() {
-  if (isLocked_) {
-    LOG(WARNING)
-        << "LazyShadowTreeRevisionConsistencyManager::lockRevisions() called without unlocking a previous lock";
-    return;
-  }
-
   // 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_) {
+  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) {
+    std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
+    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