Merge promoted React revision immediately when on the main thread (#56448)

Summary:

Changelog: [GENERAL][FIXED] When a React revision is being promoted on the main thread, merge it immediately instead of scheduling a rendering update

In the current implementation, when React is committing on the main thread, the merge is scheduled at the end of the event loop anyway. This change breaks synchronous events if they cause another render. Instead of merging the new tree immediately, it was scheduled, resulting in a visible flash of an incorrect state.

Differential Revision: D100942119

Author: j-piasecki

packages/react-native/React/Fabric/RCTScheduler.h

@@ -33,6 +33,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)schedulerShouldMergeReactRevision:(facebook::react::SurfaceId)surfaceId;
 
+- (BOOL)schedulerShouldPromoteReactRevision:(facebook::react::SurfaceId)surfaceId;
+
 - (void)schedulerDidDispatchCommand:(const facebook::react::ShadowView &)shadowView
                         commandName:(const std::string &)commandName
                                args:(const folly::dynamic &)args;

packages/react-native/React/Fabric/RCTScheduler.mm

@@ -43,6 +43,12 @@ void schedulerShouldMergeReactRevision(SurfaceId surfaceId) override
     [scheduler.delegate schedulerShouldMergeReactRevision:surfaceId];
   }
 
+  bool schedulerShouldPromoteReactRevision(SurfaceId surfaceId) override
+  {
+    RCTScheduler *scheduler = (__bridge RCTScheduler *)scheduler_;
+    return [scheduler.delegate schedulerShouldPromoteReactRevision:surfaceId];
+  }
+
   void schedulerDidRequestPreliminaryViewAllocation(const ShadowNode &shadowNode) override
   {
     // Does nothing.

packages/react-native/React/Fabric/RCTSurfacePresenter.mm

@@ -316,6 +316,20 @@ - (void)schedulerShouldMergeReactRevision:(SurfaceId)surfaceId
   });
 }
 
+- (BOOL)schedulerShouldPromoteReactRevision:(SurfaceId)surfaceId
+{
+  if (!RCTIsMainQueue()) {
+    return NO;
+  }
+
+  auto scheduler = [self scheduler];
+  scheduler.uiManager->getShadowTreeRegistry().visit(surfaceId, [](const ShadowTree &shadowTree) {
+    shadowTree.promoteReactRevision();
+    shadowTree.mergeReactRevision();
+  });
+  return YES;
+}
+
 - (void)schedulerDidDispatchCommand:(const ShadowView &)shadowView
                         commandName:(const std::string &)commandName
                                args:(const folly::dynamic &)args

packages/react-native/ReactAndroid/src/main/jni/react/fabric/FabricMountingManager.h

@@ -70,9 +70,9 @@ class FabricMountingManager final {
 
   void scheduleReactRevisionMerge(SurfaceId surfaceId);
 
- private:
   bool isOnMainThread();
 
+ private:
   jni::global_ref<JFabricUIManager::javaobject> javaUIManager_;
 
   std::recursive_mutex commitMutex_;

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

@@ -711,6 +711,23 @@ void FabricUIManagerBinding::schedulerShouldMergeReactRevision(
   }
 }
 
+bool FabricUIManagerBinding::schedulerShouldPromoteReactRevision(
+    SurfaceId surfaceId) {
+  std::shared_lock lock(installMutex_);
+  auto mountingManager =
+      getMountingManager("schedulerShouldPromoteReactRevision");
+  if (!mountingManager || !mountingManager->isOnMainThread()) {
+    return false;
+  }
+
+  scheduler_->getUIManager()->getShadowTreeRegistry().visit(
+      surfaceId, [](const ShadowTree& shadowTree) {
+        shadowTree.promoteReactRevision();
+        shadowTree.mergeReactRevision();
+      });
+  return true;
+}
+
 void FabricUIManagerBinding::mergeReactRevision(SurfaceId surfaceId) {
   std::shared_lock lock(installMutex_);
   scheduler_->getUIManager()->getShadowTreeRegistry().visit(

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

@@ -99,6 +99,8 @@ class FabricUIManagerBinding : public jni::HybridClass<FabricUIManagerBinding>,
 
   void schedulerShouldMergeReactRevision(SurfaceId surfaceId) override;
 
+  bool schedulerShouldPromoteReactRevision(SurfaceId surfaceId) override;
+
   void mergeReactRevision(SurfaceId surfaceId);
 
   void schedulerDidRequestPreliminaryViewAllocation(const ShadowNode &shadowNode) override;

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

@@ -377,6 +377,15 @@ void Scheduler::uiManagerShouldRemoveEventListener(
 
 void Scheduler::uiManagerDidFinishReactCommit(const ShadowTree& shadowTree) {
   auto surfaceId = shadowTree.getSurfaceId();
+
+  // If the commit is happening on the main thread, it should be merged
+  // immediately instead of being scheduled. This allows for synchronous
+  // events to be handled correctly.
+  if (delegate_ != nullptr &&
+      delegate_->schedulerShouldPromoteReactRevision(surfaceId)) {
+    return;
+  }
+
   runtimeScheduler_->scheduleRenderingUpdate(
       surfaceId, [surfaceId, uiManager = uiManager_]() {
         uiManager->getShadowTreeRegistry().visit(

packages/react-native/ReactCommon/react/renderer/scheduler/SchedulerDelegate.h

@@ -44,6 +44,15 @@ class SchedulerDelegate {
    */
   virtual void schedulerShouldMergeReactRevision(SurfaceId surfaceId) = 0;
 
+  /*
+   * Called when a React commit finishes and the JS revision needs to be
+   * promoted. If the platform is currently on the main thread, it should
+   * promote and merge the revision immediately, returning true.
+   * If not on the main thread, it should return false and the caller will
+   * schedule the promotion at the end of the event loop.
+   */
+  virtual bool schedulerShouldPromoteReactRevision(SurfaceId surfaceId) = 0;
+
   /*
    * Called right after a new ShadowNode was created.
    */

packages/react-native/ReactCxxPlatform/react/renderer/scheduler/SchedulerDelegateImpl.cpp

@@ -43,6 +43,11 @@ void SchedulerDelegateImpl::schedulerShouldMergeReactRevision(
       [](const ShadowTree& shadowTree) { shadowTree.mergeReactRevision(); });
 }
 
+bool SchedulerDelegateImpl::schedulerShouldPromoteReactRevision(
+    SurfaceId /*surfaceId*/) {
+  return false;
+}
+
 void SchedulerDelegateImpl::schedulerDidRequestPreliminaryViewAllocation(
     const ShadowNode& shadowNode) {}
 

packages/react-native/ReactCxxPlatform/react/renderer/scheduler/SchedulerDelegateImpl.h

@@ -35,6 +35,8 @@ class SchedulerDelegateImpl : public SchedulerDelegate {
 
   void schedulerShouldMergeReactRevision(SurfaceId surfaceId) override;
 
+  bool schedulerShouldPromoteReactRevision(SurfaceId surfaceId) override;
+
   void schedulerDidRequestPreliminaryViewAllocation(const ShadowNode &shadowNode) override;
 
   void schedulerDidDispatchCommand(