Always clean up taskview

When creating a TaskViewTaskController, it registers itself in
TaskViewTransitions in the constructor.
This means, even if the task is not initialized, we still need to clean
up the TaskViewTaskController.

Ignoring that task id is not present and always calling the removeTask
method in BubbleTaskView.
Updating TaskViewTaskController removeTask method to always remove the
class from TaskViewTransitions. Even if the task itself is not
initialized.

Bug: 369995920
Flag: com.android.wm.shell.enable_task_view_controller_cleanup
Test: atest BubbleTaskViewTest
Change-Id: Ie4766452c30a497a7a447c7473a85be1ddc45b0e

4 files changed, 42 insertions, 23 deletions

diff --git a/libs/WindowManager/Shell/multivalentTests/Android.bp b/libs/WindowManager/Shell/multivalentTests/Android.bp
index ee0d5bbed324..41d1b5c15369 100644
--- a/libs/WindowManager/Shell/multivalentTests/Android.bp
+++ b/libs/WindowManager/Shell/multivalentTests/Android.bp
@@ -53,6 +53,8 @@ android_robolectric_test {
         "mockito-robolectric-prebuilt",
         "mockito-kotlin2",
         "truth",
+        "flag-junit-base",
+        "flag-junit",
     ],
     auto_gen_config: true,
 }
diff --git a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleTaskViewTest.kt b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleTaskViewTest.kt
index 398fd554f030..4214e6ff9eaf 100644
--- a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleTaskViewTest.kt
+++ b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleTaskViewTest.kt
@@ -18,14 +18,19 @@ package com.android.wm.shell.bubbles
 
 import android.content.ComponentName
 import android.content.Context
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.wm.shell.Flags
 import com.android.wm.shell.taskview.TaskView
 
 import com.google.common.truth.Truth.assertThat
 import com.google.common.util.concurrent.MoreExecutors.directExecutor
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.kotlin.mock
@@ -36,6 +41,9 @@ import org.mockito.kotlin.verify
 @RunWith(AndroidJUnit4::class)
 class BubbleTaskViewTest {
 
+    @get:Rule
+    val setFlagsRule = SetFlagsRule()
+
     private lateinit var bubbleTaskView: BubbleTaskView
     private val context = ApplicationProvider.getApplicationContext<Context>()
     private lateinit var taskView: TaskView
@@ -75,14 +83,33 @@ class BubbleTaskViewTest {
         assertThat(actualComponentName).isEqualTo(componentName)
     }
 
+    @DisableFlags(Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP)
     @Test
-    fun cleanup_invalidTaskId_doesNotRemoveTask() {
+    fun cleanup_flagOff_invalidTaskId_doesNotRemoveTask() {
         bubbleTaskView.cleanup()
         verify(taskView, never()).removeTask()
     }
 
+    @EnableFlags(Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP)
+    @Test
+    fun cleanup_flagOn_invalidTaskId_removesTask() {
+        bubbleTaskView.cleanup()
+        verify(taskView).removeTask()
+    }
+
+    @DisableFlags(Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP)
+    @Test
+    fun cleanup_flagOff_validTaskId_removesTask() {
+        val componentName = ComponentName(context, "TestClass")
+        bubbleTaskView.listener.onTaskCreated(123, componentName)
+
+        bubbleTaskView.cleanup()
+        verify(taskView).removeTask()
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP)
     @Test
-    fun cleanup_validTaskId_removesTask() {
+    fun cleanup_flagOn_validTaskId_removesTask() {
         val componentName = ComponentName(context, "TestClass")
         bubbleTaskView.listener.onTaskCreated(123, componentName)
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskView.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskView.kt
index 65f8e48eb822..68fc0c99abee 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskView.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskView.kt
@@ -16,14 +16,11 @@
 
 package com.android.wm.shell.bubbles
 
-import android.app.ActivityTaskManager
 import android.app.ActivityTaskManager.INVALID_TASK_ID
 import android.content.ComponentName
-import android.os.RemoteException
-import android.util.Log
 import androidx.annotation.VisibleForTesting
+import com.android.wm.shell.Flags
 import com.android.wm.shell.taskview.TaskView
-import com.android.wm.shell.transition.Transitions.ENABLE_SHELL_TRANSITIONS
 import java.util.concurrent.Executor
 
 /**
@@ -89,21 +86,8 @@ class BubbleTaskView(val taskView: TaskView, executor: Executor) {
      * This should be called after all other cleanup animations have finished.
      */
     fun cleanup() {
-        if (taskId != INVALID_TASK_ID) {
-            // Ensure the task is removed from WM
-            if (ENABLE_SHELL_TRANSITIONS) {
-                taskView.removeTask()
-            } else {
-                try {
-                    ActivityTaskManager.getService().removeTask(taskId)
-                } catch (e: RemoteException) {
-                    Log.w(TAG, e.message ?: "")
-                }
-            }
+        if (Flags.enableTaskViewControllerCleanup() || taskId != INVALID_TASK_ID) {
+            taskView.removeTask()
         }
     }
-
-    private companion object {
-        const val TAG = "BubbleTaskView"
-    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
index e74342e1910c..eaeedba8ef9c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
@@ -40,6 +40,7 @@ import android.view.WindowInsets;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.SyncTransactionQueue;
 
@@ -525,8 +526,13 @@ public class TaskViewTaskController implements ShellTaskOrganizer.TaskListener {
      */
     void removeTask() {
         if (mTaskToken == null) {
-            // Call to remove task before we have one, do nothing
-            Slog.w(TAG, "Trying to remove a task that was never added? (no taskToken)");
+            if (Flags.enableTaskViewControllerCleanup()) {
+                // We don't have a task yet. Only clean up the controller
+                mTaskViewTransitions.removeTaskView(this);
+            } else {
+                // Call to remove task before we have one, do nothing
+                Slog.w(TAG, "Trying to remove a task that was never added? (no taskToken)");
+            }
             return;
         }
         mShellExecutor.execute(() -> {