feat(android): add unstable_Immediate UpdateMode for synchronous state updates

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

@@ -16,6 +16,7 @@ import com.facebook.react.bridge.ReadableNativeMap
 import com.facebook.react.bridge.WritableMap
 import com.facebook.react.common.mapbuffer.ReadableMapBuffer
 import com.facebook.react.uimanager.ReferenceStateWrapper
+import com.facebook.react.uimanager.StateWrapper
 
 /**
  * This class holds reference to the C++ EventEmitter object. Instances of this class are created on
@@ -33,7 +34,7 @@ internal class StateWrapperImpl private constructor() : HybridClassBase(), Refer
 
   private external fun getStateDataReferenceImpl(): Any?
 
-  public external fun updateStateImpl(map: NativeMap)
+  public external fun updateStateImpl(map: NativeMap, updateMode: Int)
 
   public override val stateDataMapBuffer: ReadableMapBuffer?
     get() {
@@ -66,12 +67,12 @@ internal class StateWrapperImpl private constructor() : HybridClassBase(), Refer
     initHybrid()
   }
 
-  override fun updateState(map: WritableMap) {
+  override fun updateState(map: WritableMap, updateMode: StateWrapper.UpdateMode) {
     if (!isValid) {
       FLog.e(TAG, "Race between StateWrapperImpl destruction and updateState")
       return
     }
-    updateStateImpl(map as NativeMap)
+    updateStateImpl(map as NativeMap, updateMode.value)
   }
 
   override fun destroyState() {

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/fabric/mounting/MountItemDispatcher.kt

@@ -35,6 +35,7 @@ internal class MountItemDispatcher(
   private val preMountItems: Queue<MountItem> = ConcurrentLinkedQueue()
 
   private var inDispatch: Boolean = false
+  private var followUpDispatchRequired: Boolean = false
   var batchedExecutionTime: Long = 0L
     private set
 
@@ -78,24 +79,27 @@ internal class MountItemDispatcher(
   @UiThread
   @ThreadConfined(UI)
   fun tryDispatchMountItems() {
-    // If we're already dispatching, don't reenter.
-    // Reentrance can potentially happen a lot on Android in Fabric because `updateState` from the
-    // mounting layer causes mount items to be dispatched synchronously. We want to 1) make sure we
-    // don't reenter in those cases, but 2) still execute those queued instructions synchronously.
-    // This is a pretty blunt tool, but we might not have better options since we really don't want
-    // to execute anything out-of-order.
+    // If we're already dispatching, don't reenter but signal that a follow-up dispatch is
+    // needed. This follows the same pattern as iOS's RCTMountingManager::initiateTransaction,
+    // which uses _followUpTransactionRequired flag to ensure mount items
+    // enqueued during dispatch (e.g., from synchronous state updates triggered by view layout)
+    // are processed in the same frame rather than deferred to the next one.
     if (inDispatch) {
+      followUpDispatchRequired = true
       return
     }
 
-    inDispatch = true
+    do {
+      followUpDispatchRequired = false
+      inDispatch = true
 
-    try {
-      dispatchMountItems()
-    } finally {
-      // Clean up after running dispatchMountItems - even if an exception was thrown
-      inDispatch = false
-    }
+      try {
+        dispatchMountItems()
+      } finally {
+        // Clean up after running dispatchMountItems - even if an exception was thrown
+        inDispatch = false
+      }
+    } while (followUpDispatchRequired)
 
     // We call didDispatchMountItems regardless of whether we actually dispatched anything, since
     // NativeAnimatedModule relies on this for executing any animations that may have been

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/uimanager/StateWrapper.kt

@@ -17,6 +17,15 @@ import com.facebook.react.common.mapbuffer.ReadableMapBuffer
  * by calling updateState, which communicates state back to the C++ layer.
  */
 public interface StateWrapper {
+  /**
+   * Maps to EventQueue::UpdateMode in C++.
+   * Controls how state updates are flushed (Async or Sync).
+   */
+  public enum class UpdateMode(public val value: Int) {
+    Asynchronous(0),
+    unstable_Immediate(1)
+  }
+
   /**
    * Get a ReadableMapBuffer object from the C++ layer, which is a K/V map of short keys to values.
    *
@@ -32,8 +41,9 @@ public interface StateWrapper {
   /**
    * Pass a map of values back to the C++ layer. The operation is performed synchronously and cannot
    * fail.
+   * updateMode controls whether the update is queued asynchronously or flushed immediately.
    */
-  public fun updateState(map: WritableMap)
+  public fun updateState(map: WritableMap, updateMode: UpdateMode = UpdateMode.Asynchronous)
 
   /**
    * Mark state as unused and clean up in Java and in native. This should be called as early as

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

@@ -9,6 +9,7 @@
 #include <fbjni/fbjni.h>
 #include <react/featureflags/ReactNativeFeatureFlags.h>
 #include <react/jni/ReadableNativeMap.h>
+#include <react/renderer/core/EventQueue.h>
 #include <react/renderer/mapbuffer/MapBuffer.h>
 #include <react/renderer/mapbuffer/MapBufferBuilder.h>
 
@@ -50,12 +51,14 @@ jni::local_ref<jobject> StateWrapperImpl::getStateDataReferenceImpl() {
   return nullptr;
 }
 
-void StateWrapperImpl::updateStateImpl(NativeMap* map) {
+void StateWrapperImpl::updateStateImpl(NativeMap* map, jint updateMode) {
   if (state_) {
     // Get folly::dynamic from map
     auto dynamicMap = map->consume();
     // Set state
-    state_->updateState(std::move(dynamicMap));
+    state_->updateState(
+        std::move(dynamicMap),
+        static_cast<EventQueue::UpdateMode>(updateMode));
   }
 }
 

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

@@ -27,7 +27,7 @@ class StateWrapperImpl : public jni::HybridClass<StateWrapperImpl, StateWrapper>
   jni::local_ref<JReadableMapBuffer::jhybridobject> getStateMapBufferDataImpl();
   jni::local_ref<ReadableNativeMap::jhybridobject> getStateDataImpl();
   jni::local_ref<jobject> getStateDataReferenceImpl();
-  void updateStateImpl(NativeMap *map);
+  void updateStateImpl(NativeMap *map, jint updateMode);
   void setState(std::shared_ptr<const State> state);
   std::shared_ptr<const State> getState() const;
 

packages/react-native/ReactAndroid/src/test/java/com/facebook/react/fabric/mounting/MountItemDispatcherTest.kt

@@ -0,0 +1,150 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+package com.facebook.react.fabric.mounting
+
+import com.facebook.react.fabric.mounting.MountItemDispatcher.ItemDispatchListener
+import com.facebook.react.fabric.mounting.mountitems.MountItem
+import com.facebook.react.internal.featureflags.ReactNativeFeatureFlagsForTests
+import com.facebook.react.internal.tracing.PerformanceTracer
+import com.facebook.react.uimanager.ViewManagerRegistry
+import com.facebook.testutils.shadows.ShadowSoLoader
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.annotation.Config
+import org.robolectric.annotation.Implementation
+import org.robolectric.annotation.Implements
+
+@RunWith(RobolectricTestRunner::class)
+@Config(shadows = [ShadowSoLoader::class, MountItemDispatcherTest.ShadowPerformanceTracer::class])
+class MountItemDispatcherTest {
+
+  private lateinit var dispatcher: MountItemDispatcher
+  private lateinit var dispatchListener: RecordingItemDispatchListener
+
+  @Before
+  fun setUp() {
+    ReactNativeFeatureFlagsForTests.setUp()
+    val mountingManager =
+        MountingManager(ViewManagerRegistry(emptyList()), MountingManager.MountItemExecutor {})
+    dispatchListener = RecordingItemDispatchListener()
+    dispatcher = MountItemDispatcher(mountingManager, dispatchListener)
+  }
+
+  @Test
+  fun tryDispatchMountItems_executesQueuedMountItem() {
+    val item = RecordingMountItem()
+
+    dispatcher.addMountItem(item)
+    dispatcher.tryDispatchMountItems()
+
+    assertThat(item.executed).isTrue()
+    assertThat(dispatchListener.didDispatchCount).isEqualTo(1)
+  }
+
+  /**
+   * Regression test for synchronous (`unstable_Immediate`) state updates on Android. When a mount
+   * item's execution synchronously updates shadow node state, it can enqueue new mount items and
+   * re-enter [MountItemDispatcher.tryDispatchMountItems] while a dispatch is already in progress.
+   *
+   * The re-entrant call must not drop those items: they should be flushed in the same dispatch pass,
+   * rather than being deferred to the next frame.
+   */
+  @Test
+  fun tryDispatchMountItems_reentrantDispatch_executesFollowUpItemInSamePass() {
+    val followUpItem = RecordingMountItem()
+    // Simulates a synchronous state update triggered while the first item is being mounted: it
+    // enqueues another mount item and re-enters the dispatcher.
+    val initialItem =
+        RecordingMountItem(
+            onExecute = {
+              dispatcher.addMountItem(followUpItem)
+              dispatcher.tryDispatchMountItems()
+            })
+
+    dispatcher.addMountItem(initialItem)
+    dispatcher.tryDispatchMountItems()
+
+    assertThat(initialItem.executed).isTrue()
+    assertThat(followUpItem.executed).isTrue()
+  }
+
+  @Test
+  fun tryDispatchMountItems_reentrantDispatch_preservesExecutionOrder() {
+    val executionOrder = mutableListOf<String>()
+    val followUpItem = RecordingMountItem(onExecute = { executionOrder.add("followUp") })
+    val initialItem =
+        RecordingMountItem(
+            onExecute = {
+              executionOrder.add("initial")
+              dispatcher.addMountItem(followUpItem)
+              dispatcher.tryDispatchMountItems()
+            })
+
+    dispatcher.addMountItem(initialItem)
+    dispatcher.tryDispatchMountItems()
+
+    assertThat(executionOrder).containsExactly("initial", "followUp")
+  }
+
+  @Test
+  fun tryDispatchMountItems_reentrantDispatch_invokesDidDispatchOnceForOuterCall() {
+    val initialItem =
+        RecordingMountItem(
+            onExecute = {
+              dispatcher.addMountItem(RecordingMountItem())
+              dispatcher.tryDispatchMountItems()
+            })
+
+    dispatcher.addMountItem(initialItem)
+    dispatcher.tryDispatchMountItems()
+
+    // The re-entrant call returns early and must not notify the listener; only the outer call does,
+    // once, after the follow-up loop has drained everything.
+    assertThat(dispatchListener.didDispatchCount).isEqualTo(1)
+  }
+
+  private class RecordingMountItem(
+      private val surfaceId: Int = 1,
+      private val onExecute: () -> Unit = {},
+  ) : MountItem {
+    var executed: Boolean = false
+      private set
+
+    override fun execute(mountingManager: MountingManager) {
+      executed = true
+      onExecute()
+    }
+
+    override fun getSurfaceId(): Int = surfaceId
+  }
+
+  private class RecordingItemDispatchListener : ItemDispatchListener {
+    var didDispatchCount: Int = 0
+      private set
+
+    override fun willMountItems(mountItems: List<MountItem>?) = Unit
+
+    override fun didMountItems(mountItems: List<MountItem>?) = Unit
+
+    override fun didDispatchMountItems() {
+      didDispatchCount++
+    }
+  }
+
+  // isTracing() is a native method; its JNI library isn't loaded in JVM tests, so return false and
+  // let trace() just run its block. ShadowSoLoader covers the SoLoader.loadLibrary init.
+  @Implements(PerformanceTracer::class)
+  class ShadowPerformanceTracer {
+    companion object {
+      @JvmStatic @Implementation fun isTracing(): Boolean = false
+    }
+  }
+}

packages/react-native/ReactAndroid/src/test/java/com/facebook/react/uimanager/StateWrapperUpdateModeTest.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+package com.facebook.react.uimanager
+
+import com.facebook.react.uimanager.StateWrapper.UpdateMode
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+/**
+ * UpdateMode.value crosses JNI and is cast straight to C++ EventQueue::UpdateMode, so the values
+ * must match its order (Asynchronous = 0, unstable_Immediate = 1). A mismatch silently flips
+ * async/sync state updates.
+ */
+@RunWith(RobolectricTestRunner::class)
+class StateWrapperUpdateModeTest {
+
+  @Test
+  fun updateModeValues_matchCxxEventQueueUpdateMode() {
+    assertThat(UpdateMode.Asynchronous.value).isEqualTo(0)
+    assertThat(UpdateMode.unstable_Immediate.value).isEqualTo(1)
+  }
+
+  @Test
+  fun updateModeValues_areDistinct() {
+    assertThat(UpdateMode.entries.map { it.value }).doesNotHaveDuplicates()
+  }
+}

packages/react-native/ReactCommon/react/renderer/core/ConcreteState.h

@@ -116,9 +116,11 @@ class ConcreteState : public State {
     return getData().getDynamic();
   }
 
-  void updateState(folly::dynamic &&data) const override
+  void updateState(
+      folly::dynamic &&data,
+      EventQueue::UpdateMode updateMode = EventQueue::UpdateMode::Asynchronous) const override
   {
-    updateState(Data(getData(), std::move(data)));
+    updateState(Data(getData(), std::move(data)), updateMode);
   }
 
   MapBuffer getMapBuffer() const override

packages/react-native/ReactCommon/react/renderer/core/State.h

@@ -10,6 +10,7 @@
 #ifdef RN_SERIALIZABLE_STATE
 #include <fbjni/fbjni.h>
 #include <folly/dynamic.h>
+#include <react/renderer/core/EventQueue.h>
 #include <react/renderer/mapbuffer/MapBuffer.h>
 #endif
 
@@ -66,7 +67,9 @@ class State {
   virtual folly::dynamic getDynamic() const = 0;
   virtual MapBuffer getMapBuffer() const = 0;
   virtual jni::local_ref<jobject> getJNIReference() const = 0;
-  virtual void updateState(folly::dynamic &&data) const = 0;
+  virtual void updateState(
+      folly::dynamic &&data,
+      EventQueue::UpdateMode updateMode = EventQueue::UpdateMode::Asynchronous) const = 0;
 #endif
 
  protected: