synchronous updateState on android

Author: intergalacticspacehighway

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 a do-while loop with a _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/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: