Fix SurfaceMountingManager leaking views from stopped surfaces

Summary:
When a Surface is stopped, we don't immediately destroy the SurfaceMountingManager but instead just tear down its internal state. This allows for better error handling (eg did this react tag ever exist, or is this non-existing tag).

The way we construct the set of tags post-deletion is flawed though: `mTagToViewState.keySet()` does not create a new Set with all the tags used, but instead uses the underlying HashMap to iterate over the keys as needed. This effectively keeps all the Views inside that deleted surface alive.

Changelog: [Android][Fixed] Surfaces in the new architecture no longer leak views once stopped

Differential Revision: D46840717

fbshipit-source-id: ead356ba3ee8d0b1f6e22fbf848f36655e3c4472

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/config/ReactFeatureFlags.java

@@ -159,4 +159,7 @@ public class ReactFeatureFlags {
 
   /** Report mount operations from the host platform to notify mount hooks. */
   public static boolean enableMountHooks = false;
+
+  /** Fixes a leak in SurfaceMountingManager.mTagSetForStoppedSurface */
+  public static boolean fixStoppedSurfaceTagSetLeak = true;
 }

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

@@ -16,6 +16,7 @@
 import androidx.annotation.AnyThread;
 import androidx.annotation.NonNull;
 import androidx.annotation.UiThread;
+import androidx.collection.SparseArrayCompat;
 import com.facebook.common.logging.FLog;
 import com.facebook.infer.annotation.Assertions;
 import com.facebook.infer.annotation.ThreadConfined;
@@ -52,6 +53,7 @@
 import com.facebook.react.views.view.ReactViewManagerWrapper;
 import java.util.HashSet;
 import java.util.LinkedList;
+import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
 import java.util.Stack;
@@ -91,7 +93,8 @@ public class SurfaceMountingManager {
   private RemoveDeleteTreeUIFrameCallback mRemoveDeleteTreeUIFrameCallback;
 
   // This is null *until* StopSurface is called.
-  private Set<Integer> mTagSetForStoppedSurface;
+  private Set<Integer> mTagSetForStoppedSurfaceLegacy;
+  private SparseArrayCompat<Object> mTagSetForStoppedSurface;
 
   private final int mSurfaceId;
 
@@ -166,7 +169,10 @@ public boolean getViewExists(int tag) {
     // If Surface stopped, check if tag *was* associated with this Surface, even though it's been
     // deleted. This helps distinguish between scenarios where an invalid tag is referenced, vs
     // race conditions where an imperative method is called on a tag during/just after StopSurface.
-    if (mTagSetForStoppedSurface != null && mTagSetForStoppedSurface.contains(tag)) {
+    if (mTagSetForStoppedSurface != null && mTagSetForStoppedSurface.containsKey(tag)) {
+      return true;
+    }
+    if (mTagSetForStoppedSurfaceLegacy != null && mTagSetForStoppedSurfaceLegacy.contains(tag)) {
       return true;
     }
     if (mTagToViewState == null) {
@@ -287,27 +293,38 @@ public void stopSurface() {
     }
 
     Runnable runnable =
-        new Runnable() {
-          @Override
-          public void run() {
-            // We must call `onDropViewInstance` on all remaining Views
+        () -> {
+          if (ReactFeatureFlags.fixStoppedSurfaceTagSetLeak) {
+            mTagSetForStoppedSurface = new SparseArrayCompat<>();
+            for (Map.Entry<Integer, ViewState> entry : mTagToViewState.entrySet()) {
+              // Using this as a placeholder value in the map. We're using SparseArrayCompat
+              // since it can efficiently represent the list of pending tags
+              mTagSetForStoppedSurface.put(entry.getKey(), this);
+
+              // We must call `onDropViewInstance` on all remaining Views
+              onViewStateDeleted(entry.getValue());
+            }
+          } else {
             for (ViewState viewState : mTagToViewState.values()) {
+              // We must call `onDropViewInstance` on all remaining Views
               onViewStateDeleted(viewState);
             }
+            mTagSetForStoppedSurfaceLegacy = mTagToViewState.keySet();
+          }
 
-            // Evict all views from cache and memory
-            mTagSetForStoppedSurface = mTagToViewState.keySet();
-            mTagToViewState = null;
-            mJSResponderHandler = null;
-            mRootViewManager = null;
-            mMountItemExecutor = null;
-            mOnViewAttachItems.clear();
-
-            if (ReactFeatureFlags.enableViewRecycling) {
-              mViewManagerRegistry.onSurfaceStopped(mSurfaceId);
-            }
-            FLog.e(TAG, "Surface [" + mSurfaceId + "] was stopped on SurfaceMountingManager.");
+          // Evict all views from cache and memory
+          // TODO: clear instead of nulling out to simplify null-safety in this class
+          mTagToViewState = null;
+          mJSResponderHandler = null;
+          mRootViewManager = null;
+          mMountItemExecutor = null;
+          mThemedReactContext = null;
+          mOnViewAttachItems.clear();
+
+          if (ReactFeatureFlags.enableViewRecycling) {
+            mViewManagerRegistry.onSurfaceStopped(mSurfaceId);
           }
+          FLog.e(TAG, "Surface [" + mSurfaceId + "] was stopped on SurfaceMountingManager.");
         };
 
     if (UiThreadUtil.isOnUiThread()) {