Fix potential race conditions when executing commitHooks

packages/react-native/ReactCommon/react/renderer/imagemanager/platform/android/react/renderer/imagemanager/ImageFetcher.cpp

@@ -7,36 +7,57 @@
 
 #include "ImageFetcher.h"
 
-#include <glog/logging.h>
 #include <react/common/mapbuffer/JReadableMapBuffer.h>
 #include <react/renderer/imagemanager/conversions.h>
+#include <react/renderer/runtimescheduler/RuntimeScheduler.h>
+#include <react/renderer/uimanager/UIManager.h>
+#include <react/renderer/uimanager/UIManagerBinding.h>
 
 namespace facebook::react {
 
-ImageFetcher::ImageFetcher(
-    std::shared_ptr<const ContextContainer> contextContainer)
-    : contextContainer_(std::move(contextContainer)) {
-  if (ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
-    if (auto uiManagerCommitHookManager =
-            contextContainer_
-                ->find<std::shared_ptr<UIManagerCommitHookManager>>(
-                    std::string(UIManagerCommitHookManagerKey));
-        uiManagerCommitHookManager.has_value()) {
-      (*uiManagerCommitHookManager)->registerCommitHook(*this);
-    }
+using RunOnUIManagerFn = std::function<void(UIManager&, ImageFetcher&)>;
+
+namespace {
+
+void runOnUIManager(
+    const ContextContainer& contextContainer,
+    std::shared_ptr<ImageFetcher> strongThis,
+    RunOnUIManagerFn&& runOnUIManagerFn) {
+  auto weakRuntimeScheduler =
+      contextContainer.find<std::weak_ptr<RuntimeScheduler>>(
+          RuntimeSchedulerKey);
+  auto runtimeScheduler = weakRuntimeScheduler.has_value()
+      ? weakRuntimeScheduler.value().lock()
+      : nullptr;
+  if (runtimeScheduler == nullptr) {
+    return;
   }
+  runtimeScheduler->scheduleTask(
+      SchedulerPriority::NormalPriority,
+      [strongThis,
+       runOnUIManagerFn = std::move(runOnUIManagerFn)](jsi::Runtime& runtime) {
+        runOnUIManagerFn(
+            UIManagerBinding::getBinding(runtime)->getUIManager(), *strongThis);
+      });
 }
 
+} // namespace
+
+ImageFetcher::ImageFetcher(
+    std::shared_ptr<const ContextContainer> contextContainer)
+    : contextContainer_(std::move(contextContainer)) {}
+
 ImageFetcher::~ImageFetcher() {
-  if (ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
-    if (auto uiManagerCommitHookManager =
-            contextContainer_
-                ->find<std::shared_ptr<UIManagerCommitHookManager>>(
-                    std::string(UIManagerCommitHookManagerKey));
-        uiManagerCommitHookManager.has_value()) {
-      (*uiManagerCommitHookManager)->unregisterCommitHook(*this);
-    }
+  if (!commitHookRegistered_) {
+    return;
   }
+  commitHookRegistered_ = false;
+  runOnUIManager(
+      *contextContainer_,
+      shared_from_this(),
+      [](UIManager& uiManager, ImageFetcher& strongThis) {
+        uiManager.unregisterCommitHook(strongThis);
+      });
 }
 
 ImageRequest ImageFetcher::requestImage(
@@ -51,7 +72,17 @@ ImageRequest ImageFetcher::requestImage(
 
   auto telemetry = std::make_shared<ImageTelemetry>(surfaceId);
 
-  if (!ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
+  if (ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
+    if (!commitHookRegistered_) {
+      runOnUIManager(
+          *contextContainer_,
+          shared_from_this(),
+          [](UIManager& uiManager, ImageFetcher& strongThis) {
+            uiManager.registerCommitHook(strongThis);
+          });
+      commitHookRegistered_ = true;
+    }
+  } else {
     flushImageRequests();
   }
 
@@ -63,7 +94,9 @@ RootShadowNode::Unshared ImageFetcher::shadowTreeWillCommit(
     const RootShadowNode::Shared& /*oldRootShadowNode*/,
     const RootShadowNode::Unshared& newRootShadowNode,
     const ShadowTree::CommitOptions& /*commitOptions*/) noexcept {
-  flushImageRequests();
+  if (commitHookRegistered_) {
+    flushImageRequests();
+  }
   return newRootShadowNode;
 }
 

packages/react-native/ReactCommon/react/renderer/imagemanager/platform/android/react/renderer/imagemanager/ImageFetcher.h

@@ -18,7 +18,8 @@
 
 namespace facebook::react {
 
-class ImageFetcher : public UIManagerCommitHook {
+class ImageFetcher : public UIManagerCommitHook,
+                     public std::enable_shared_from_this<ImageFetcher> {
  public:
   ImageFetcher(std::shared_ptr<const ContextContainer> contextContainer);
   ~ImageFetcher() override;
@@ -47,7 +48,9 @@ class ImageFetcher : public UIManagerCommitHook {
  private:
   void flushImageRequests();
 
+  bool commitHookRegistered_{false};
   std::unordered_map<SurfaceId, std::vector<ImageRequestItem>> items_;
   std::shared_ptr<const ContextContainer> contextContainer_;
 };
+
 } // namespace facebook::react

packages/react-native/ReactCommon/react/renderer/imagemanager/platform/android/react/renderer/imagemanager/ImageManager.cpp

@@ -18,14 +18,22 @@ constexpr inline bool isInteger(const std::string& str) {
   return str.find_first_not_of("0123456789") == std::string::npos;
 }
 
+struct ImageFetcherHolder {
+  explicit ImageFetcherHolder(
+      const std::shared_ptr<const ContextContainer>& contextContainer)
+      : imageFetcher_(std::make_shared<ImageFetcher>(contextContainer)) {}
+
+  const std::shared_ptr<ImageFetcher> imageFetcher_;
+};
+
 } // namespace
 
 ImageManager::ImageManager(
     const std::shared_ptr<const ContextContainer>& contextContainer)
-    : self_(new ImageFetcher(contextContainer)) {}
+    : self_(new ImageFetcherHolder(contextContainer)) {}
 
 ImageManager::~ImageManager() {
-  delete static_cast<ImageFetcher*>(self_);
+  delete static_cast<ImageFetcherHolder*>(self_);
 }
 
 ImageRequest ImageManager::requestImage(
@@ -35,8 +43,9 @@ ImageRequest ImageManager::requestImage(
     Tag tag) const {
   if (ReactNativeFeatureFlags::enableImagePrefetchingAndroid()) {
     if (!isInteger(imageSource.uri)) {
-      return static_cast<ImageFetcher*>(self_)->requestImage(
-          imageSource, surfaceId, imageRequestParams, tag);
+      return static_cast<ImageFetcherHolder*>(self_)
+          ->imageFetcher_->requestImage(
+              imageSource, surfaceId, imageRequestParams, tag);
     }
   }
   return {imageSource, nullptr, {}};

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

@@ -24,33 +24,6 @@
 
 namespace facebook::react {
 
-namespace {
-
-class UIManagerCommitHookManagerImpl : public UIManagerCommitHookManager {
- public:
-  explicit UIManagerCommitHookManagerImpl(std::weak_ptr<UIManager> uiManager)
-      : uiManager_(std::move(uiManager)) {}
-
-  ~UIManagerCommitHookManagerImpl() override = default;
-
-  void registerCommitHook(UIManagerCommitHook& commitHook) override {
-    if (auto uiManager = uiManager_.lock()) {
-      uiManager->registerCommitHook(commitHook);
-    }
-  }
-
-  void unregisterCommitHook(UIManagerCommitHook& commitHook) override {
-    if (auto uiManager = uiManager_.lock()) {
-      uiManager->unregisterCommitHook(commitHook);
-    }
-  }
-
- private:
-  std::weak_ptr<UIManager> uiManager_;
-};
-
-} // namespace
-
 Scheduler::Scheduler(
     const SchedulerToolbox& schedulerToolbox,
     UIManagerAnimationDelegate* animationDelegate,
@@ -76,11 +49,6 @@ Scheduler::Scheduler(
 
   auto uiManager =
       std::make_shared<UIManager>(runtimeExecutor_, contextContainer_);
-  if (ReactNativeFeatureFlags::enableImagePrefetchingAndroid()) {
-    contextContainer_->insert(
-        std::string(UIManagerCommitHookManagerKey),
-        std::make_shared<UIManagerCommitHookManagerImpl>(uiManager));
-  }
 
   auto eventOwnerBox = std::make_shared<EventBeat::OwnerBox>();
   eventOwnerBox->owner = eventDispatcher_;

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

@@ -57,15 +57,4 @@ class UIManagerCommitHook {
   virtual ~UIManagerCommitHook() noexcept = default;
 };
 
-constexpr std::string_view UIManagerCommitHookManagerKey =
-    "UIManagerCommitHookManager";
-
-class UIManagerCommitHookManager {
- public:
-  virtual ~UIManagerCommitHookManager() = default;
-
-  virtual void registerCommitHook(UIManagerCommitHook& commitHook) = 0;
-  virtual void unregisterCommitHook(UIManagerCommitHook& commitHook) = 0;
-};
-
 } // namespace facebook::react