Use a weakref hashmap to store controllers

Protects against memory leaks. TaskViewTaskController adds itself to the
map when the object is created.

Bug: 369995920
Test: atest TaskViewTransitionsTest
Test: revert change to fix memory leak with BubbleTaskView cleanup,
  trigger a bubble and swipe to dismiss without opening it, observe from
  dump that TaskViewTransitions is holding on to the TaskViewContorller
  for that task, trigger a gc for systemui, observe that after gc the
  TaskViewController is cleared up from TaskViewTransitions

Flag: com.android.wm.shell.enable_task_view_controller_cleanup
Change-Id: I180f432e6de5200b63ab17c09504b8ecc32a8292

3 files changed, 54 insertions, 8 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
index 39648f65b4f3..ae75cb589de0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
@@ -37,11 +37,14 @@ import android.window.WindowContainerTransaction;
 
 import androidx.annotation.VisibleForTesting;
 
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.shared.TransitionUtil;
 import com.android.wm.shell.transition.Transitions;
 
 import java.util.ArrayList;
+import java.util.Map;
 import java.util.Objects;
+import java.util.WeakHashMap;
 
 /**
  * Handles Shell Transitions that involve TaskView tasks.
@@ -49,8 +52,15 @@ import java.util.Objects;
 public class TaskViewTransitions implements Transitions.TransitionHandler {
     static final String TAG = "TaskViewTransitions";
 
-    private final ArrayMap<TaskViewTaskController, TaskViewRequestedState> mTaskViews =
-            new ArrayMap<>();
+    /**
+     * Map of {@link TaskViewTaskController} to {@link TaskViewRequestedState}.
+     * <p>
+     * {@link TaskView} keeps a reference to the {@link TaskViewTaskController} instance and
+     * manages its lifecycle.
+     * Only keep a weak reference to the controller instance here to allow for it to be cleaned
+     * up when its TaskView is no longer used.
+     */
+    private final Map<TaskViewTaskController, TaskViewRequestedState> mTaskViews;
     private final ArrayList<PendingTransition> mPending = new ArrayList<>();
     private final Transitions mTransitions;
     private final boolean[] mRegistered = new boolean[]{false};
@@ -95,6 +105,11 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
 
     public TaskViewTransitions(Transitions transitions) {
         mTransitions = transitions;
+        if (Flags.enableTaskViewControllerCleanup()) {
+            mTaskViews = new WeakHashMap<>();
+        } else {
+            mTaskViews = new ArrayMap<>();
+        }
         // Defer registration until the first TaskView because we want this to be the "first" in
         // priority when handling requests.
         // TODO(210041388): register here once we have an explicit ordering mechanism.
@@ -208,10 +223,21 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
     }
 
     private TaskViewTaskController findTaskView(ActivityManager.RunningTaskInfo taskInfo) {
-        for (int i = 0; i < mTaskViews.size(); ++i) {
-            if (mTaskViews.keyAt(i).getTaskInfo() == null) continue;
-            if (taskInfo.token.equals(mTaskViews.keyAt(i).getTaskInfo().token)) {
-                return mTaskViews.keyAt(i);
+        if (Flags.enableTaskViewControllerCleanup()) {
+            for (TaskViewTaskController controller : mTaskViews.keySet()) {
+                if (controller.getTaskInfo() == null) continue;
+                if (taskInfo.token.equals(controller.getTaskInfo().token)) {
+                    return controller;
+                }
+            }
+        } else {
+            ArrayMap<TaskViewTaskController, TaskViewRequestedState> taskViews =
+                    (ArrayMap<TaskViewTaskController, TaskViewRequestedState>) mTaskViews;
+            for (int i = 0; i < taskViews.size(); ++i) {
+                if (taskViews.keyAt(i).getTaskInfo() == null) continue;
+                if (taskInfo.token.equals(taskViews.keyAt(i).getTaskInfo().token)) {
+                    return taskViews.keyAt(i);
+                }
             }
         }
         return null;
diff --git a/libs/WindowManager/Shell/tests/unittest/Android.bp b/libs/WindowManager/Shell/tests/unittest/Android.bp
index 049a5a0615e0..91be5f58b1f7 100644
--- a/libs/WindowManager/Shell/tests/unittest/Android.bp
+++ b/libs/WindowManager/Shell/tests/unittest/Android.bp
@@ -63,6 +63,7 @@ android_test {
         "com.android.window.flags.window-aconfig-java",
         "platform-test-annotations",
         "flag-junit",
+        "platform-parametric-runner-lib",
     ],
 
     libs: [
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
index d3e40f21db23..60e203009e8c 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
@@ -33,7 +33,8 @@ import static org.mockito.Mockito.when;
 import android.app.ActivityManager;
 import android.graphics.Rect;
 import android.os.IBinder;
-import android.testing.AndroidTestingRunner;
+import android.platform.test.flag.junit.FlagsParameterization;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.TestableLooper;
 import android.view.SurfaceControl;
 import android.window.TransitionInfo;
@@ -42,10 +43,12 @@ import android.window.WindowContainerTransaction;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.transition.Transitions;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -54,11 +57,23 @@ import org.mockito.MockitoAnnotations;
 import java.util.ArrayList;
 import java.util.List;
 
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
+import platform.test.runner.parameterized.Parameters;
+
 @SmallTest
-@RunWith(AndroidTestingRunner.class)
+@RunWith(ParameterizedAndroidJunit4.class)
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
 public class TaskViewTransitionsTest extends ShellTestCase {
 
+    @Parameters(name = "{0}")
+    public static List<FlagsParameterization> getParams() {
+        return FlagsParameterization.allCombinationsOf(
+                Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP);
+    }
+
+    @Rule
+    public final SetFlagsRule mSetFlagsRule;
+
     @Mock
     Transitions mTransitions;
     @Mock
@@ -70,6 +85,10 @@ public class TaskViewTransitionsTest extends ShellTestCase {
 
     TaskViewTransitions mTaskViewTransitions;
 
+    public TaskViewTransitionsTest(FlagsParameterization flags) {
+        mSetFlagsRule = new SetFlagsRule(flags);
+    }
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);