Fix DispatchUIFrameCallback invoked multiple times per frame (#46346)

Summary:
Pull Request resolved: #46346

Noticed that we could sometimes have multiple instances of DispatchUIFrameCallback in a single frame, which would cause us to execute more view preallocations or other work scheduled by Fabric. Root cause for this is on the new architecture, onHostResume seems to be invoked multiple times.

Make this code more resilient by explicitly tracking the state of the frame callback and avoiding multiple subscriptions. Longer-term we should consider having ReactChoreographer support repeating FrameCallbacks, since most of them are.

Changelog: [Android][Fixed] Fixed multiple Fabric dispatch callbacks being executed in a single Android frame

Reviewed By: sammy-SC

Differential Revision: D62213721

fbshipit-source-id: ac6fa5483ea38d9a15824af233fd23f1f6f3c891

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/fabric/FabricUIManager.java

@@ -421,11 +421,6 @@ public void invalidate() {
 
     mDestroyed = true;
 
-    // This is not technically thread-safe, since it's read on the UI thread and written
-    // here on the JS thread. We've marked it as volatile so that this writes to UI-thread
-    // memory immediately.
-    mDispatchUIFrameCallback.stop();
-
     mEventDispatcher.removeBatchEventDispatchedListener(mBatchEventDispatchedListener);
     mEventDispatcher.unregisterEventEmitter(FABRIC);
 
@@ -1021,8 +1016,7 @@ public void receiveEvent(
 
   @Override
   public void onHostResume() {
-    ReactChoreographer.getInstance()
-        .postFrameCallback(ReactChoreographer.CallbackType.DISPATCH_UI, mDispatchUIFrameCallback);
+    mDispatchUIFrameCallback.resume();
   }
 
   @Override
@@ -1034,8 +1028,7 @@ public EventDispatcher getEventDispatcher() {
 
   @Override
   public void onHostPause() {
-    ReactChoreographer.getInstance()
-        .removeFrameCallback(ReactChoreographer.CallbackType.DISPATCH_UI, mDispatchUIFrameCallback);
+    mDispatchUIFrameCallback.pause();
   }
 
   @Override
@@ -1322,21 +1315,55 @@ private class DispatchUIFrameCallback extends GuardedFrameCallback {
 
     private volatile boolean mIsMountingEnabled = true;
 
+    @ThreadConfined(UI)
+    private boolean mShouldSchedule = false;
+
+    @ThreadConfined(UI)
+    private boolean mIsScheduled = false;
+
     private DispatchUIFrameCallback(@NonNull ReactContext reactContext) {
       super(reactContext);
     }
 
-    @AnyThread
-    void stop() {
-      mIsMountingEnabled = false;
+    @UiThread
+    @ThreadConfined(UI)
+    private void schedule() {
+      if (!mIsScheduled && mShouldSchedule) {
+        mIsScheduled = true;
+        ReactChoreographer.getInstance()
+            .postFrameCallback(ReactChoreographer.CallbackType.DISPATCH_UI, this);
+      }
+    }
+
+    @UiThread
+    @ThreadConfined(UI)
+    void resume() {
+      mShouldSchedule = true;
+      schedule();
+    }
+
+    @UiThread
+    @ThreadConfined(UI)
+    void pause() {
+      ReactChoreographer.getInstance()
+          .removeFrameCallback(ReactChoreographer.CallbackType.DISPATCH_UI, this);
+      mShouldSchedule = false;
+      mIsScheduled = false;
     }
 
     @Override
     @UiThread
     @ThreadConfined(UI)
     public void doFrameGuarded(long frameTimeNanos) {
-      if (!mIsMountingEnabled || mDestroyed) {
-        FLog.w(TAG, "Not flushing pending UI operations because of previously thrown Exception");
+      mIsScheduled = false;
+
+      if (!mIsMountingEnabled) {
+        FLog.w(TAG, "Not flushing pending UI operations: exception was previously thrown");
+        return;
+      }
+
+      if (mDestroyed) {
+        FLog.w(TAG, "Not flushing pending UI operations: FabricUIManager is destroyed");
         return;
       }
 
@@ -1362,12 +1389,10 @@ public void doFrameGuarded(long frameTimeNanos) {
         mMountItemDispatcher.tryDispatchMountItems();
       } catch (Exception ex) {
         FLog.e(TAG, "Exception thrown when executing UIFrameGuarded", ex);
-        stop();
+        mIsMountingEnabled = false;
         throw ex;
       } finally {
-        ReactChoreographer.getInstance()
-            .postFrameCallback(
-                ReactChoreographer.CallbackType.DISPATCH_UI, mDispatchUIFrameCallback);
+        schedule();
       }
 
       mSynchronousEvents.clear();

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/modules/core/ReactChoreographer.kt

@@ -7,6 +7,7 @@
 
 package com.facebook.react.modules.core
 
+import androidx.annotation.GuardedBy
 import android.view.Choreographer
 import com.facebook.common.logging.FLog
 import com.facebook.infer.annotation.Assertions
@@ -41,12 +42,12 @@ public class ReactChoreographer private constructor(choreographerProvider: Chore
   private val callbackQueues: Array<ArrayDeque<Choreographer.FrameCallback>> =
       Array(CallbackType.entries.size) { ArrayDeque() }
   private var totalCallbacks = 0
+  @GuardedBy("callbackQueues")
   private var hasPostedCallback = false
 
   private val frameCallback =
       Choreographer.FrameCallback { frameTimeNanos ->
         synchronized(callbackQueues) {
-
           // Callbacks run once and are then automatically removed, the callback will
           // be posted again from postFrameCallback
           hasPostedCallback = false
@@ -76,17 +77,7 @@ public class ReactChoreographer private constructor(choreographerProvider: Chore
       callbackQueues[type.order].addLast(callback)
       totalCallbacks++
       Assertions.assertCondition(totalCallbacks > 0)
-      if (!hasPostedCallback) {
-        if (choreographer == null) {
-          // Schedule on the main thread, at which point the constructor's async work will have
-          // completed
-          UiThreadUtil.runOnUiThread {
-            synchronized(callbackQueues) { postFrameCallbackOnChoreographer() }
-          }
-        } else {
-          postFrameCallbackOnChoreographer()
-        }
-      }
+      postFrameCallbackOnChoreographer()
     }
   }
 
@@ -102,12 +93,22 @@ public class ReactChoreographer private constructor(choreographerProvider: Chore
   }
 
   /**
-   * This method writes on mHasPostedCallback and it should be called from another method that has
+   * This method writes [hasPostedCallback] and it should be called from another method that has
    * the lock on [callbackQueues].
    */
   private fun postFrameCallbackOnChoreographer() {
-    choreographer?.postFrameCallback(frameCallback)
-    hasPostedCallback = true
+    if (!hasPostedCallback) {
+      val choreographer = choreographer
+      if (choreographer == null) {
+        // Schedule on the main thread, at which point the constructor's async work will have
+         UiThreadUtil.runOnUiThread {
+                   synchronized(callbackQueues) { postFrameCallbackOnChoreographer() }
+                  }
+      } else {
+        choreographer.postFrameCallback(frameCallback)
+        hasPostedCallback = true
+      }
+    }
   }
 
   /**