Fabric: Introducing ShadowTreeRegistry

Summary:
Why do we need a dedicated registry class?
* We need to simplify registry-related logic in Scheduler.
* We need to couple threading aspect of the registry with the registry itself, otherwise it's not clear why exactly we acquire the mutex. We also should not acquire the mutex in a per-method way (as we did before), because it's incorrect and misleading (only lines that access the registry should by protected).
* We need to have a way to share the registry with other classes (e.g. UIManager) without passing a reference to the whole Scheduler.

Reviewed By: mdvacca

Differential Revision: D13036550

fbshipit-source-id: 644da910e823666c586834a3da2b4cdcb90eebb2

ReactCommon/fabric/uimanager/Scheduler.cpp

@@ -66,12 +66,11 @@ void Scheduler::startSurface(
     const folly::dynamic &initialProps,
     const LayoutConstraints &layoutConstraints,
     const LayoutContext &layoutContext) const {
-  std::lock_guard<std::mutex> lock(mutex_);
-
   auto shadowTree =
       std::make_unique<ShadowTree>(surfaceId, layoutConstraints, layoutContext);
   shadowTree->setDelegate(this);
-  shadowTreeRegistry_.emplace(surfaceId, std::move(shadowTree));
+
+  shadowTreeRegistry_.add(std::move(shadowTree));
 
 #ifndef ANDROID
   runtimeExecutor_([=](jsi::Runtime &runtime) {
@@ -96,11 +95,10 @@ void Scheduler::renderTemplateToSurface(
         *componentDescriptorRegistry_,
         nMR);
 
-    std::lock_guard<std::mutex> lock(mutex_);
-    const auto &shadowTree = shadowTreeRegistry_.at(surfaceId);
-    assert(shadowTree);
-    shadowTree->complete(
-        std::make_shared<SharedShadowNodeList>(SharedShadowNodeList{tree}));
+    shadowTreeRegistry_.get(surfaceId, [=](const ShadowTree &shadowTree) {
+      shadowTree.complete(
+          std::make_shared<SharedShadowNodeList>(SharedShadowNodeList{tree}));
+    });
   } catch (const std::exception &e) {
     LOG(ERROR) << "    >>>> EXCEPTION <<<  rendering uiTemplate in "
                << "Scheduler::renderTemplateToSurface: " << e.what();
@@ -108,14 +106,14 @@ void Scheduler::renderTemplateToSurface(
 }
 
 void Scheduler::stopSurface(SurfaceId surfaceId) const {
-  std::lock_guard<std::mutex> lock(mutex_);
-  const auto &iterator = shadowTreeRegistry_.find(surfaceId);
-  auto &shadowTree = *iterator->second;
-  // As part of stopping the Surface, we have to commit an empty tree.
-  shadowTree.complete(std::const_pointer_cast<SharedShadowNodeList>(
-      ShadowNode::emptySharedShadowNodeSharedList()));
-  shadowTree.setDelegate(nullptr);
-  shadowTreeRegistry_.erase(iterator);
+  shadowTreeRegistry_.get(surfaceId, [](const ShadowTree &shadowTree) {
+    // As part of stopping the Surface, we have to commit an empty tree.
+    shadowTree.complete(std::const_pointer_cast<SharedShadowNodeList>(
+        ShadowNode::emptySharedShadowNodeSharedList()));
+  });
+
+  auto shadowTree = shadowTreeRegistry_.remove(surfaceId);
+  shadowTree->setDelegate(nullptr);
 
 #ifndef ANDROID
   runtimeExecutor_([=](jsi::Runtime &runtime) {
@@ -128,21 +126,21 @@ Size Scheduler::measureSurface(
     SurfaceId surfaceId,
     const LayoutConstraints &layoutConstraints,
     const LayoutContext &layoutContext) const {
-  std::lock_guard<std::mutex> lock(mutex_);
-  const auto &shadowTree = shadowTreeRegistry_.at(surfaceId);
-  assert(shadowTree);
-  return shadowTree->measure(layoutConstraints, layoutContext);
+  Size size;
+  shadowTreeRegistry_.get(surfaceId, [&](const ShadowTree &shadowTree) {
+    size = shadowTree.measure(layoutConstraints, layoutContext);
+  });
+  return size;
 }
 
 void Scheduler::constraintSurfaceLayout(
     SurfaceId surfaceId,
     const LayoutConstraints &layoutConstraints,
     const LayoutContext &layoutContext) const {
-  std::lock_guard<std::mutex> lock(mutex_);
-  const auto &shadowTree = shadowTreeRegistry_.at(surfaceId);
-  assert(shadowTree);
-  shadowTree->synchronize([&]() {
-    shadowTree->constraintLayout(layoutConstraints, layoutContext);
+  shadowTreeRegistry_.get(surfaceId, [&](const ShadowTree &shadowTree) {
+    shadowTree.synchronize([&]() {
+      shadowTree.constraintLayout(layoutConstraints, layoutContext);
+    });
   });
 }
 
@@ -170,16 +168,11 @@ void Scheduler::shadowTreeDidCommit(
 #pragma mark - UIManagerDelegate
 
 void Scheduler::uiManagerDidFinishTransaction(
-    Tag rootTag,
+    SurfaceId surfaceId,
     const SharedShadowNodeUnsharedList &rootChildNodes) {
-  std::lock_guard<std::mutex> lock(mutex_);
-  const auto iterator = shadowTreeRegistry_.find(rootTag);
-  if (iterator == shadowTreeRegistry_.end()) {
-    // This might happen during surface unmounting/deallocation process
-    // due to the asynchronous nature of JS calls.
-    return;
-  }
-  iterator->second->complete(rootChildNodes);
+  shadowTreeRegistry_.get(surfaceId, [&](const ShadowTree &shadowTree) {
+    shadowTree.complete(rootChildNodes);
+  });
 }
 
 void Scheduler::uiManagerDidCreateShadowNode(

ReactCommon/fabric/uimanager/Scheduler.h

@@ -15,6 +15,7 @@
 #include <react/uimanager/SchedulerDelegate.h>
 #include <react/uimanager/ShadowTree.h>
 #include <react/uimanager/ShadowTreeDelegate.h>
+#include <react/uimanager/ShadowTreeRegistry.h>
 #include <react/uimanager/UIManagerBinding.h>
 #include <react/uimanager/UIManagerDelegate.h>
 #include <react/uimanager/primitives.h>
@@ -75,7 +76,7 @@ class Scheduler final : public UIManagerDelegate, public ShadowTreeDelegate {
 #pragma mark - UIManagerDelegate
 
   void uiManagerDidFinishTransaction(
-      Tag rootTag,
+      SurfaceId surfaceId,
       const SharedShadowNodeUnsharedList &rootChildNodes) override;
   void uiManagerDidCreateShadowNode(
       const SharedShadowNode &shadowNode) override;
@@ -89,9 +90,7 @@ class Scheduler final : public UIManagerDelegate, public ShadowTreeDelegate {
  private:
   SchedulerDelegate *delegate_;
   SharedComponentDescriptorRegistry componentDescriptorRegistry_;
-  mutable std::mutex mutex_;
-  mutable std::unordered_map<SurfaceId, std::unique_ptr<ShadowTree>>
-      shadowTreeRegistry_; // Protected by `mutex_`.
+  ShadowTreeRegistry shadowTreeRegistry_;
   SharedContextContainer contextContainer_;
   RuntimeExecutor runtimeExecutor_;
   std::shared_ptr<UIManagerBinding> uiManagerBinding_;

ReactCommon/fabric/uimanager/ShadowTreeRegistry.cpp

@@ -0,0 +1,40 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "ShadowTreeRegistry.h"
+
+namespace facebook {
+namespace react {
+
+void ShadowTreeRegistry::add(std::unique_ptr<ShadowTree> &&shadowTree) const {
+  std::lock_guard<std::mutex> lock(mutex_);
+  registry_.emplace(shadowTree->getSurfaceId(), std::move(shadowTree));
+}
+
+std::unique_ptr<ShadowTree> ShadowTreeRegistry::remove(
+    SurfaceId surfaceId) const {
+  std::lock_guard<std::mutex> lock(mutex_);
+  auto iterator = registry_.find(surfaceId);
+  auto shadowTree = std::unique_ptr<ShadowTree>(iterator->second.release());
+  registry_.erase(iterator);
+  return shadowTree;
+}
+
+bool ShadowTreeRegistry::get(
+    SurfaceId surfaceId,
+    std::function<void(const ShadowTree &shadowTree)> callback) const {
+  std::lock_guard<std::mutex> lock(mutex_);
+  auto iterator = registry_.find(surfaceId);
+
+  if (iterator == registry_.end()) {
+    return false;
+  }
+
+  callback(*iterator->second);
+  return true;
+}
+
+} // namespace react
+} // namespace facebook

ReactCommon/fabric/uimanager/ShadowTreeRegistry.h

@@ -0,0 +1,55 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <react/core/ReactPrimitives.h>
+#include <react/uimanager/ShadowTree.h>
+
+namespace facebook {
+namespace react {
+
+/*
+ * Owning registry of `ShadowTree`s.
+ */
+class ShadowTreeRegistry final {
+ public:
+  ShadowTreeRegistry() = default;
+
+  /*
+   * Adds a `ShadowTree` instance to the registry.
+   * The ownership of the instance is also transferred to the registry.
+   * Can be called from any thread.
+   */
+  void add(std::unique_ptr<ShadowTree> &&shadowTree) const;
+
+  /*
+   * Removes a `ShadowTree` instance with given `surfaceId` from the registry
+   * and returns it as a result.
+   * The ownership of the instance is also transferred to the caller.
+   * Can be called from any thread.
+   */
+  std::unique_ptr<ShadowTree> remove(SurfaceId surfaceId) const;
+
+  /*
+   * Finds a `ShadowTree` instance with a given `surfaceId` in the registry and
+   * synchronously calls the `callback` with a reference to the instance while
+   * the mutex is being acquired.
+   * Returns `true` if the registry has `ShadowTree` instance with corresponding
+   * `surfaceId`, otherwise returns `false` without calling the `callback`.
+   * Can be called from any thread.
+   */
+  bool get(
+      SurfaceId surfaceId,
+      std::function<void(const ShadowTree &shadowTree)> callback) const;
+
+ private:
+  mutable std::mutex mutex_;
+  mutable std::unordered_map<SurfaceId, std::unique_ptr<ShadowTree>>
+      registry_; // Protected by `mutex_`.
+};
+
+} // namespace react
+} // namespace facebook

ReactCommon/fabric/uimanager/UIManagerDelegate.h

@@ -23,7 +23,7 @@ class UIManagerDelegate {
    * The tree is not layed out and not sealed at this time.
    */
   virtual void uiManagerDidFinishTransaction(
-      Tag rootTag,
+      SurfaceId surfaceId,
       const SharedShadowNodeUnsharedList &rootChildNodes) = 0;
 
   /*