Image Prefetching > Re-enable JNI batching for ImageFetcher (#54486)

Summary:

This brings back another version of batching image request JNI calls. Previous attempts crashed in edge cases

Changelog: [Internal]

Reviewed By: lenaic

Differential Revision: D86676643

Author: christophpurrer

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>
@@ -636,6 +637,18 @@ void FabricUIManagerBinding::schedulerDidFinishTransaction(
 
 void FabricUIManagerBinding::schedulerShouldRenderTransactions(
     const std::shared_ptr<const MountingCoordinator>& mountingCoordinator) {
+  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();
+    }
+  }
+
   auto mountingManager =
       getMountingManager("schedulerShouldRenderTransactions");
   if (!mountingManager) {

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

@@ -8,14 +8,23 @@
 #include "ImageFetcher.h"
 
 #include <react/common/mapbuffer/JReadableMapBuffer.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)) {}
 
+ImageFetcher::~ImageFetcher() noexcept {
+  if (ReactNativeFeatureFlags::enableImagePrefetchingJNIBatchingAndroid()) {
+    contextContainer_->erase(ImageFetcherKey);
+  }
+}
+
 ImageRequest ImageFetcher::requestImage(
     const ImageSource& imageSource,
     SurfaceId surfaceId,
@@ -29,9 +38,11 @@ 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() {

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

@@ -16,24 +16,27 @@
 
 namespace facebook::react {
 
+extern const char ImageFetcherKey[];
+
 class ImageFetcher {
  public:
   ImageFetcher(std::shared_ptr<const ContextContainer> contextContainer);
-  ~ImageFetcher() = default;
+  ~ImageFetcher() noexcept;
   ImageFetcher(const ImageFetcher &) = delete;
   ImageFetcher &operator=(const ImageFetcher &) = delete;
   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_;
 };

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

@@ -22,10 +22,17 @@ constexpr inline bool isInteger(const std::string& str) {
 
 ImageManager::ImageManager(
     const std::shared_ptr<const ContextContainer>& contextContainer)
-    : self_(new ImageFetcher(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_);
+  delete static_cast<std::shared_ptr<ImageFetcher>*>(self_);
 }
 
 ImageRequest ImageManager::requestImage(
@@ -35,8 +42,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, {}};