Destroy Fabric State objects earlier in Java

Summary:
There are races between BackgroundExecutor and Fabric/View teardown, where, because of the way things are set up currently, a View will strongly retain a pointer into some native State object which can keep a host of other objects alive in C++, even after stopSurface, and even after RN itself starts tearing down.

To alleviate this, we more aggressively clear State from Java, without waiting for Java GC: 1) on stopSurface, 2) whenever a State object is stale from Java's perspective.

This should allow us to keep all common updateState semantics, while only introducing a new edge-case that stateWrapper can be destroyed during mounting if stopSurface happens at the same time. In those cases, checking for NPEs should be sufficient.

The possible race condition only really happens with updateState, so it's easier to check for. There should practically be no cases where there's a race between `stopSurface` and `getState`, because `getState` is only really called around state updates and view preallocation, and never after; we still check for NPEs in those cases, but it shouldn't be an issue.

Changelog: [Internal]

Reviewed By: thurn, sammy-SC, mdvacca

Differential Revision: D26858584

fbshipit-source-id: 2ef7467220865380037d69d8de322fe8797f6a12

ReactAndroid/src/main/java/com/facebook/react/fabric/StateWrapperImpl.java

@@ -9,6 +9,8 @@ package com.facebook.react.fabric;
 
 import android.annotation.SuppressLint;
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import com.facebook.common.logging.FLog;
 import com.facebook.jni.HybridData;
 import com.facebook.proguard.annotations.DoNotStrip;
 import com.facebook.react.bridge.NativeMap;
@@ -26,7 +28,10 @@ public class StateWrapperImpl implements StateWrapper {
     FabricSoLoader.staticInit();
   }
 
+  private static final String TAG = "StateWrapperImpl";
+
   @DoNotStrip private final HybridData mHybridData;
+  private volatile boolean mDestroyed = false;
 
   private static native HybridData initHybrid();
 
@@ -34,13 +39,34 @@ public class StateWrapperImpl implements StateWrapper {
     mHybridData = initHybrid();
   }
 
+  private native ReadableNativeMap getStateDataImpl();
+
   @Override
-  public native ReadableNativeMap getState();
+  @Nullable
+  public ReadableNativeMap getStateData() {
+    if (mDestroyed) {
+      FLog.e(TAG, "Race between StateWrapperImpl destruction and getState");
+      return null;
+    }
+    return getStateDataImpl();
+  }
 
   public native void updateStateImpl(@NonNull NativeMap map);
 
   @Override
   public void updateState(@NonNull WritableMap map) {
+    if (mDestroyed) {
+      FLog.e(TAG, "Race between StateWrapperImpl destruction and updateState");
+      return;
+    }
     updateStateImpl((NativeMap) map);
   }
+
+  @Override
+  public void destroyState() {
+    if (!mDestroyed) {
+      mDestroyed = true;
+      mHybridData.resetNative();
+    }
+  }
 }

ReactAndroid/src/main/java/com/facebook/react/fabric/jni/StateWrapperImpl.cpp

@@ -22,7 +22,8 @@ jni::local_ref<StateWrapperImpl::jhybriddata> StateWrapperImpl::initHybrid(
   return makeCxxInstance();
 }
 
-jni::local_ref<ReadableNativeMap::jhybridobject> StateWrapperImpl::getState() {
+jni::local_ref<ReadableNativeMap::jhybridobject>
+StateWrapperImpl::getStateDataImpl() {
   folly::dynamic map = state_->getDynamic();
   local_ref<ReadableNativeMap::jhybridobject> readableNativeMap =
       ReadableNativeMap::newObjectCxxArgs(map);
@@ -39,7 +40,7 @@ void StateWrapperImpl::updateStateImpl(NativeMap *map) {
 void StateWrapperImpl::registerNatives() {
   registerHybrid({
       makeNativeMethod("initHybrid", StateWrapperImpl::initHybrid),
-      makeNativeMethod("getState", StateWrapperImpl::getState),
+      makeNativeMethod("getStateDataImpl", StateWrapperImpl::getStateDataImpl),
       makeNativeMethod("updateStateImpl", StateWrapperImpl::updateStateImpl),
   });
 }

ReactAndroid/src/main/java/com/facebook/react/fabric/jni/StateWrapperImpl.h

@@ -25,7 +25,7 @@ class StateWrapperImpl : public jni::HybridClass<StateWrapperImpl> {
 
   static void registerNatives();
 
-  jni::local_ref<ReadableNativeMap::jhybridobject> getState();
+  jni::local_ref<ReadableNativeMap::jhybridobject> getStateDataImpl();
   void updateStateImpl(NativeMap *map);
   void updateStateWithFailureCallbackImpl(
       NativeMap *map,

ReactAndroid/src/main/java/com/facebook/react/fabric/mounting/SurfaceMountingManager.java

@@ -21,7 +21,6 @@ import com.facebook.infer.annotation.ThreadConfined;
 import com.facebook.react.bridge.ReactSoftException;
 import com.facebook.react.bridge.ReadableArray;
 import com.facebook.react.bridge.ReadableMap;
-import com.facebook.react.bridge.ReadableNativeMap;
 import com.facebook.react.bridge.RetryableMountingLayerException;
 import com.facebook.react.bridge.SoftAssertions;
 import com.facebook.react.bridge.UiThreadUtil;
@@ -204,9 +203,20 @@ public class SurfaceMountingManager {
       return;
     }
 
-    // Prevent more views from being created, or the hierarchy from being manipulated at all
+    // Prevent more views from being created, or the hierarchy from being manipulated at all. This
+    // causes further operations to noop.
     mIsStopped = true;
 
+    // Reset all StateWrapper objects
+    // Since this can happen on any thread, is it possible to race between StateWrapper destruction
+    // and some accesses from View classes in the UI thread?
+    for (ViewState viewState : mTagToViewState.values()) {
+      if (viewState.mStateWrapper != null) {
+        viewState.mStateWrapper.destroyState();
+        viewState.mStateWrapper = null;
+      }
+    }
+
     Runnable runnable =
         new Runnable() {
           @Override
@@ -502,7 +512,7 @@ public class SurfaceMountingManager {
 
     ViewState viewState = new ViewState(reactTag, view, viewManager);
     viewState.mCurrentProps = propsDiffMap;
-    viewState.mCurrentState = (stateWrapper != null ? stateWrapper.getState() : null);
+    viewState.mStateWrapper = stateWrapper;
 
     mTagToViewState.put(reactTag, viewState);
   }
@@ -675,9 +685,9 @@ public class SurfaceMountingManager {
     }
 
     ViewState viewState = getViewState(reactTag);
-    @Nullable ReadableNativeMap newState = stateWrapper == null ? null : stateWrapper.getState();
 
-    viewState.mCurrentState = newState;
+    StateWrapper prevStateWrapper = viewState.mStateWrapper;
+    viewState.mStateWrapper = stateWrapper;
 
     ViewManager viewManager = viewState.mViewManager;
 
@@ -689,6 +699,12 @@ public class SurfaceMountingManager {
     if (extraData != null) {
       viewManager.updateExtraData(viewState.mView, extraData);
     }
+
+    // Immediately clear native side of previous state wrapper. This causes the State object in C++
+    // to be destroyed immediately instead of waiting for Java GC to kick in.
+    if (prevStateWrapper != null) {
+      prevStateWrapper.destroyState();
+    }
   }
 
   @UiThread
@@ -832,7 +848,7 @@ public class SurfaceMountingManager {
     @Nullable final ViewManager mViewManager;
     @Nullable public ReactStylesDiffMap mCurrentProps = null;
     @Nullable public ReadableMap mCurrentLocalData = null;
-    @Nullable public ReadableMap mCurrentState = null;
+    @Nullable public StateWrapper mStateWrapper = null;
     @Nullable public EventEmitterWrapper mEventEmitter = null;
 
     private ViewState(int reactTag, @Nullable View view, @Nullable ViewManager viewManager) {

ReactAndroid/src/main/java/com/facebook/react/uimanager/FabricViewStateManager.java

@@ -88,7 +88,7 @@ public class FabricViewStateManager {
     setState(mStateWrapper, stateUpdateCallback, 0);
   }
 
-  public @Nullable ReadableMap getState() {
-    return mStateWrapper != null ? mStateWrapper.getState() : null;
+  public @Nullable ReadableMap getStateData() {
+    return mStateWrapper != null ? mStateWrapper.getStateData() : null;
   }
 }

ReactAndroid/src/main/java/com/facebook/react/uimanager/StateWrapper.java

@@ -9,6 +9,7 @@ package com.facebook.react.uimanager;
 
 import com.facebook.react.bridge.ReadableNativeMap;
 import com.facebook.react.bridge.WritableMap;
+import javax.annotation.Nullable;
 
 /**
  * This is a wrapper that can be used for passing State objects from Fabric C++ core to
@@ -19,11 +20,20 @@ public interface StateWrapper {
   /**
    * Get a ReadableNativeMap object from the C++ layer, which is a K/V map of string keys to values.
    */
-  ReadableNativeMap getState();
+  @Nullable
+  ReadableNativeMap getStateData();
 
   /**
    * Pass a map of values back to the C++ layer. The operation is performed synchronously and cannot
    * fail.
    */
   void updateState(WritableMap map);
+
+  /**
+   * Mark state as unused and clean up in Java and in native. This should be called as early as
+   * possible when you know a StateWrapper will no longer be used. If there's ANY chance of it being
+   * used legitimately, don't destroy it! It is expected that all StateWrappers are destroyed
+   * immediately upon stopSurface.
+   */
+  void destroyState();
 }

ReactAndroid/src/main/java/com/facebook/react/views/modal/ReactModalHostView.java

@@ -471,7 +471,7 @@ public class ReactModalHostView extends ViewGroup
 
       // Check incoming state values. If they're already the correct value, return early to prevent
       // infinite UpdateState/SetState loop.
-      ReadableMap currentState = getFabricViewStateManager().getState();
+      ReadableMap currentState = getFabricViewStateManager().getStateData();
       if (currentState != null) {
         float delta = (float) 0.9;
         float stateScreenHeight =

ReactAndroid/src/main/java/com/facebook/react/views/text/ReactTextViewManager.java

@@ -83,7 +83,15 @@ public class ReactTextViewManager
   @Override
   public Object updateState(
       ReactTextView view, ReactStylesDiffMap props, @Nullable StateWrapper stateWrapper) {
-    ReadableNativeMap state = stateWrapper.getState();
+    if (stateWrapper == null) {
+      return null;
+    }
+
+    ReadableNativeMap state = stateWrapper.getStateData();
+    if (state == null) {
+      return null;
+    }
+
     ReadableMap attributedString = state.getMap("attributedString");
     ReadableMap paragraphAttributes = state.getMap("paragraphAttributes");
     Spannable spanned =

ReactAndroid/src/main/java/com/facebook/react/views/textinput/ReactTextInputManager.java

@@ -1210,7 +1210,15 @@ public class ReactTextInputManager extends BaseViewManager<ReactEditText, Layout
 
     view.getFabricViewStateManager().setStateWrapper(stateWrapper);
 
-    ReadableNativeMap state = stateWrapper.getState();
+    if (stateWrapper == null) {
+      return null;
+    }
+
+    ReadableNativeMap state = stateWrapper.getStateData();
+
+    if (state == null) {
+      return null;
+    }
 
     if (!state.hasKey("attributedString")) {
       return null;