Asset ImageFetcher thread correctness in DEBUG builds

packages/react-native/ReactAndroid/src/main/jni/react/fabric/FabricUIManagerBinding.cpp

@@ -24,6 +24,7 @@
 #include <react/renderer/core/EventBeat.h>
 #include <react/renderer/core/EventEmitter.h>
 #include <react/renderer/core/conversions.h>
+#include <react/renderer/imagemanager/ImageFetcher.h>
 #include <react/renderer/scheduler/Scheduler.h>
 #include <react/renderer/scheduler/SchedulerDelegate.h>
 #include <react/renderer/scheduler/SchedulerToolbox.h>
@@ -641,6 +642,19 @@ void FabricUIManagerBinding::schedulerShouldRenderTransactions(
   if (!mountingManager) {
     return;
   }
+
+  if (ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
+    auto weakImageFetcher =
+        scheduler_->getContextContainer()->find<std::weak_ptr<ImageFetcher>>(
+            ImageFetcherKey);
+    auto imageFetcher = weakImageFetcher.has_value()
+        ? weakImageFetcher.value().lock()
+        : nullptr;
+    if (imageFetcher != nullptr) {
+      imageFetcher->flushImageRequests();
+    }
+  }
+
   if (ReactNativeFeatureFlags::enableAccumulatedUpdatesInRawPropsAndroid()) {
     auto mountingTransaction = mountingCoordinator->pullTransaction(
         /* willPerformAsynchronously = */ true);

packages/react-native/ReactCommon/react/renderer/imagemanager/ImageManager.h

@@ -37,6 +37,9 @@ class ImageManager {
       Tag tag = {}) const;
 
  private:
+#ifdef ANDROID
+  std::shared_ptr<const ContextContainer> contextContainer_{};
+#endif
   void *self_{};
 };
 

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

@@ -8,13 +8,23 @@
 #include "ImageFetcher.h"
 
 #include <react/common/mapbuffer/JReadableMapBuffer.h>
+#include <react/debug/react_native_assert.h>
+#include <react/featureflags/ReactNativeFeatureFlags.h>
 #include <react/renderer/imagemanager/conversions.h>
 
 namespace facebook::react {
 
+extern const char ImageFetcherKey[] = "ImageFetcher";
+
 ImageFetcher::ImageFetcher(
     std::shared_ptr<const ContextContainer> contextContainer)
-    : contextContainer_(std::move(contextContainer)) {}
+    : contextContainer_(std::move(contextContainer))
+#ifdef REACT_NATIVE_DEBUG
+      ,
+      threadId_(std::this_thread::get_id())
+#endif
+{
+}
 
 ImageRequest ImageFetcher::requestImage(
     const ImageSource& imageSource,
@@ -29,12 +39,20 @@ ImageRequest ImageFetcher::requestImage(
 
   auto telemetry = std::make_shared<ImageTelemetry>(surfaceId);
 
-  flushImageRequests();
+  if (!ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
+    flushImageRequests();
+  }
 
-  return {imageSource, telemetry};
+  return ImageRequest{imageSource, telemetry};
 }
 
 void ImageFetcher::flushImageRequests() {
+#ifdef REACT_NATIVE_DEBUG
+  if (ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
+    assertThread();
+  }
+#endif
+
   if (items_.empty()) {
     return;
   }
@@ -57,4 +75,12 @@ void ImageFetcher::flushImageRequests() {
   items_.clear();
 }
 
+#ifdef REACT_NATIVE_DEBUG
+void ImageFetcher::assertThread() const {
+  react_native_assert(
+      threadId_ != std::this_thread::get_id() &&
+      "ImageFetcher method called from the wrong thread!");
+}
+#endif
+
 } // namespace facebook::react

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

@@ -13,9 +13,14 @@
 #include <memory>
 #include <unordered_map>
 #include <vector>
+#ifdef REACT_NATIVE_DEBUG
+#include <thread>
+#endif
 
 namespace facebook::react {
 
+extern const char ImageFetcherKey[];
+
 class ImageFetcher {
  public:
   ImageFetcher(std::shared_ptr<const ContextContainer> contextContainer);
@@ -25,16 +30,22 @@ class ImageFetcher {
   ImageFetcher(ImageFetcher &&) = delete;
   ImageFetcher &operator=(ImageFetcher &&) = delete;
 
+  void flushImageRequests();
+
+ private:
+  friend class ImageManager;
   ImageRequest requestImage(
       const ImageSource &imageSource,
       SurfaceId surfaceId,
       const ImageRequestParams &imageRequestParams,
       Tag tag);
 
- private:
-  void flushImageRequests();
-
   std::unordered_map<SurfaceId, std::vector<ImageRequestItem>> items_;
   std::shared_ptr<const ContextContainer> contextContainer_;
+
+#ifdef REACT_NATIVE_DEBUG
+  std::thread::id threadId_;
+  void assertThread() const;
+#endif
 };
 } // namespace facebook::react

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

@@ -22,10 +22,21 @@ constexpr inline bool isInteger(const std::string& str) {
 
 ImageManager::ImageManager(
     const std::shared_ptr<const ContextContainer>& contextContainer)
-    : self_(new ImageFetcher(contextContainer)) {}
+    : contextContainer_(contextContainer),
+      self_(new std::shared_ptr<ImageFetcher>(
+          std::make_shared<ImageFetcher>(contextContainer))) {
+  if (ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
+    std::weak_ptr<ImageFetcher> weakImageFetcher =
+        *static_cast<std::shared_ptr<ImageFetcher>*>(self_);
+    contextContainer->insert(ImageFetcherKey, weakImageFetcher);
+  }
+}
 
 ImageManager::~ImageManager() {
-  delete static_cast<ImageFetcher*>(self_);
+  if (ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
+    contextContainer_->erase(ImageFetcherKey);
+  }
+  delete static_cast<std::shared_ptr<ImageFetcher>*>(self_);
 }
 
 ImageRequest ImageManager::requestImage(
@@ -35,8 +46,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<std::shared_ptr<ImageFetcher>*>(self_)
+          ->get()
+          ->requestImage(imageSource, surfaceId, imageRequestParams, tag);
     }
   }
   return {imageSource, nullptr, {}};