Image Prefetching > Re-enable JNI batching for ImageFetcher (facebook#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>
@@ -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,10 +8,13 @@
 #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)) {}
@@ -29,9 +32,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,6 +16,8 @@
 
 namespace facebook::react {
 
+extern const char ImageFetcherKey[];
+
 class ImageFetcher {
  public:
   ImageFetcher(std::shared_ptr<const ContextContainer> contextContainer);
@@ -25,15 +27,16 @@ 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_;
 };

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, {}};