allow setting SurfaceStartedCallback on UIManager

Summary:
## Changelog:

[General] [Added] - allow setting SurfaceStartedCallback on UIManager

Reviewed By: rshest

Differential Revision: D71917124

fbshipit-source-id: b8f8eaef892518e3aaa3da82ab91b06173a644a0

Author: zeyap

packages/react-native/ReactCommon/react/renderer/scheduler/Scheduler.cpp

@@ -336,6 +336,13 @@ void Scheduler::uiManagerShouldRemoveEventListener(
   removeEventListener(listener);
 }
 
+void Scheduler::uiManagerDidStartSurface(const ShadowTree& shadowTree) {
+  std::shared_lock lock(onSurfaceStartCallbackMutex_);
+  if (onSurfaceStartCallback_) {
+    onSurfaceStartCallback_(shadowTree);
+  }
+}
+
 void Scheduler::reportMount(SurfaceId surfaceId) const {
   uiManager_->reportMount(surfaceId);
 }
@@ -362,4 +369,10 @@ void Scheduler::removeEventListener(
   }
 }
 
+void Scheduler::uiManagerShouldSetOnSurfaceStartCallback(
+    OnSurfaceStartCallback&& callback) {
+  std::shared_lock lock(onSurfaceStartCallbackMutex_);
+  onSurfaceStartCallback_ = std::move(callback);
+}
+
 } // namespace facebook::react

packages/react-native/ReactCommon/react/renderer/scheduler/Scheduler.h

@@ -101,6 +101,7 @@ class Scheduler final : public UIManagerDelegate {
       std::shared_ptr<const EventListener> listener) final;
   void uiManagerShouldRemoveEventListener(
       const std::shared_ptr<const EventListener>& listener) final;
+  void uiManagerDidStartSurface(const ShadowTree& shadowTree) override;
 
 #pragma mark - ContextContainer
   ContextContainer::Shared getContextContainer() const;
@@ -115,6 +116,10 @@ class Scheduler final : public UIManagerDelegate {
   void removeEventListener(
       const std::shared_ptr<const EventListener>& listener);
 
+#pragma mark - Surface start callback
+  void uiManagerShouldSetOnSurfaceStartCallback(
+      OnSurfaceStartCallback&& callback) override;
+
  private:
   friend class SurfaceHandler;
 
@@ -144,6 +149,9 @@ class Scheduler final : public UIManagerDelegate {
   ContextContainer::Shared contextContainer_;
 
   RuntimeScheduler* runtimeScheduler_{nullptr};
+
+  mutable std::shared_mutex onSurfaceStartCallbackMutex_;
+  OnSurfaceStartCallback onSurfaceStartCallback_;
 };
 
 } // namespace facebook::react

packages/react-native/ReactCommon/react/renderer/uimanager/UIManager.cpp

@@ -227,6 +227,13 @@ void UIManager::startSurface(
   auto surfaceId = shadowTree->getSurfaceId();
   shadowTreeRegistry_.add(std::move(shadowTree));
 
+  shadowTreeRegistry_.visit(
+      surfaceId, [delegate = delegate_](const ShadowTree& shadowTree) {
+        if (delegate != nullptr) {
+          delegate->uiManagerDidStartSurface(shadowTree);
+        }
+      });
+
   runtimeExecutor_([=](jsi::Runtime& runtime) {
     TraceSection s("UIManager::startSurface::onRuntime");
     AppRegistryBinding::startSurface(
@@ -706,4 +713,11 @@ void UIManager::removeEventListener(
   }
 }
 
+void UIManager::setOnSurfaceStartCallback(
+    UIManagerDelegate::OnSurfaceStartCallback&& callback) {
+  if (delegate_ != nullptr) {
+    delegate_->uiManagerShouldSetOnSurfaceStartCallback(std::move(callback));
+  }
+}
+
 } // namespace facebook::react

packages/react-native/ReactCommon/react/renderer/uimanager/UIManager.h

@@ -216,6 +216,10 @@ class UIManager final : public ShadowTreeDelegate {
   void removeEventListener(
       const std::shared_ptr<const EventListener>& listener);
 
+#pragma mark - Set on surface start callback
+  void setOnSurfaceStartCallback(
+      UIManagerDelegate::OnSurfaceStartCallback&& callback);
+
  private:
   friend class UIManagerBinding;
   friend class Scheduler;

packages/react-native/ReactCommon/react/renderer/uimanager/UIManagerDelegate.h

@@ -10,6 +10,7 @@
 #include <react/renderer/core/ReactPrimitives.h>
 #include <react/renderer/core/ShadowNode.h>
 #include <react/renderer/mounting/MountingCoordinator.h>
+#include <react/renderer/mounting/ShadowTree.h>
 
 namespace facebook::react {
 
@@ -77,6 +78,16 @@ class UIManagerDelegate {
   virtual void uiManagerShouldRemoveEventListener(
       const std::shared_ptr<const EventListener>& listener) = 0;
 
+  /*
+   * Start surface.
+   */
+  virtual void uiManagerDidStartSurface(const ShadowTree& shadowTree) = 0;
+
+  using OnSurfaceStartCallback =
+      std::function<void(const ShadowTree& shadowTree)>;
+  virtual void uiManagerShouldSetOnSurfaceStartCallback(
+      OnSurfaceStartCallback&& callback) = 0;
+
   virtual ~UIManagerDelegate() noexcept = default;
 };