[32/N] Desks: Use RunOnTransitStart to apply "pending" minimize/activations

This change moves both minimization (task-limit) and desk activation WCT
changes into a single, reusable util (#addDeskActivationChanges) and
updates callsites of both to use it. To accomodate both features having
to "wait" for the transition token to be available to add their "pending
state" in their respective observers, it changes the return value of the
util to a |RunOnTransitStart| callback that is invoked when the token is
ready.

This removes the burden on callers of knowing the specific side effects
of #addDeskActivationChanges and having to update the feature-specific
observers with the pending transition state.

Ideally, all other side-effects (e.g. immersive ones) should be migrated
into this method, but this CL focuses on just these two as a starting
point to unblock multi-desk development.

Flag: com.android.window.flags.enable_multiple_desktops_backend
Bug: 362720497
Test: atest WMShellUnitTests
Change-Id: I9d95813271f65ea6f58af2cf20ff2ddcea54b52c

3 files changed, 63 insertions, 81 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index b03d7396c944..b31cc1de10bf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -598,7 +598,7 @@ class DesktopTasksController(
                 reason = DesktopImmersiveController.ExitReason.TASK_LAUNCH,
             )
 
-        val taskIdToMinimize = addDeskActivationWithMovingTaskChanges(deskId, wct, task)
+        val runOnTransitStart = addDeskActivationWithMovingTaskChanges(deskId, wct, task)
 
         val transition: IBinder
         if (remoteTransition != null) {
@@ -613,20 +613,9 @@ class DesktopTasksController(
         desktopModeEnterExitTransitionListener?.onEnterDesktopModeTransitionStarted(
             FREEFORM_ANIMATION_DURATION
         )
-        taskIdToMinimize?.let {
-            addPendingMinimizeTransition(transition, it, MinimizeReason.TASK_LIMIT)
-        }
+        runOnTransitStart?.invoke(transition)
         exitResult.asExit()?.runOnTransitionStart?.invoke(transition)
-        if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-            desksTransitionObserver.addPendingTransition(
-                DeskTransition.ActiveDeskWithTask(
-                    token = transition,
-                    displayId = displayId,
-                    deskId = deskId,
-                    enterTaskId = task.taskId,
-                )
-            )
-        } else {
+        if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
             taskRepository.setActiveDesk(displayId = displayId, deskId = deskId)
         }
         return true
@@ -689,7 +678,7 @@ class DesktopTasksController(
                 moveHomeTask(context.displayId, wct)
             }
         }
-        val taskIdToMinimize = addDeskActivationWithMovingTaskChanges(deskId, wct, taskInfo)
+        val runOnTransitStart = addDeskActivationWithMovingTaskChanges(deskId, wct, taskInfo)
         val exitResult =
             desktopImmersiveController.exitImmersiveIfApplicable(
                 wct = wct,
@@ -702,20 +691,9 @@ class DesktopTasksController(
             DRAG_TO_DESKTOP_FINISH_ANIM_DURATION_MS.toInt()
         )
         if (transition != null) {
-            taskIdToMinimize?.let {
-                addPendingMinimizeTransition(transition, it, MinimizeReason.TASK_LIMIT)
-            }
+            runOnTransitStart?.invoke(transition)
             exitResult.asExit()?.runOnTransitionStart?.invoke(transition)
-            if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-                desksTransitionObserver.addPendingTransition(
-                    DeskTransition.ActiveDeskWithTask(
-                        token = transition,
-                        displayId = taskInfo.displayId,
-                        deskId = deskId,
-                        enterTaskId = taskInfo.taskId,
-                    )
-                )
-            } else {
+            if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
                 taskRepository.setActiveDesk(displayId = taskInfo.displayId, deskId = deskId)
             }
         } else {
@@ -1074,19 +1052,14 @@ class DesktopTasksController(
                         ?: getDefaultDeskId(displayId)
                 )
             val activateDeskWct = WindowContainerTransaction()
-            addDeskActivationChanges(deskIdToActivate, activateDeskWct)
+            // TODO: b/391485148 - pass in the launching task here to apply task-limit policy,
+            //  but make sure to not do it twice since it is also done at the start of this
+            //  function.
+            activationRunOnTransitStart =
+                addDeskActivationChanges(deskIdToActivate, activateDeskWct)
             // Desk activation must be handled before app launch-related transactions.
             activateDeskWct.merge(launchTransaction, /* transfer= */ true)
             launchTransaction = activateDeskWct
-            activationRunOnTransitStart = { transition ->
-                desksTransitionObserver.addPendingTransition(
-                    DeskTransition.ActivateDesk(
-                        token = transition,
-                        displayId = displayId,
-                        deskId = deskIdToActivate,
-                    )
-                )
-            }
             desktopModeEnterExitTransitionListener?.onEnterDesktopModeTransitionStarted(
                 FREEFORM_ANIMATION_DURATION
             )
@@ -1268,17 +1241,8 @@ class DesktopTasksController(
             wct.reparent(task.token, displayAreaInfo.token, /* onTop= */ true)
         }
 
-        addDeskActivationChanges(destinationDeskId, wct)
-        val activationRunnable: RunOnTransitStart = { transition ->
-            desksTransitionObserver.addPendingTransition(
-                DeskTransition.ActiveDeskWithTask(
-                    token = transition,
-                    displayId = displayId,
-                    deskId = destinationDeskId,
-                    enterTaskId = task.taskId,
-                )
-            )
-        }
+        // TODO: b/391485148 - pass in the moving-to-desk |task| here to apply task-limit policy.
+        val activationRunnable = addDeskActivationChanges(destinationDeskId, wct)
 
         if (Flags.enableDisplayFocusInShellTransitions()) {
             // Bring the destination display to top with includingParents=true, so that the
@@ -2501,10 +2465,10 @@ class DesktopTasksController(
         deskId: Int,
         wct: WindowContainerTransaction,
         task: RunningTaskInfo,
-    ): Int? {
-        val taskIdToMinimize = addDeskActivationChanges(deskId, wct, task)
+    ): RunOnTransitStart? {
+        val runOnTransitStart = addDeskActivationChanges(deskId, wct, task)
         addMoveToDeskTaskChanges(wct = wct, task = task, deskId = deskId)
-        return taskIdToMinimize
+        return runOnTransitStart
     }
 
     /**
@@ -2780,31 +2744,57 @@ class DesktopTasksController(
         deskId: Int,
         wct: WindowContainerTransaction,
         newTask: RunningTaskInfo? = null,
-    ): Int? {
+    ): RunOnTransitStart? {
+        logV("addDeskActivationChanges newTaskId=%d deskId=%d", newTask?.taskId, deskId)
+        val newTaskIdInFront = newTask?.taskId
         val displayId = taskRepository.getDisplayForDesk(deskId)
-        return if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-            prepareForDeskActivation(displayId, wct)
-            desksOrganizer.activateDesk(wct, deskId)
-            if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {
-                // TODO: 362720497 - do non-running tasks need to be restarted with |wct#startTask|?
-            }
-            taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(
-                doesAnyTaskRequireTaskbarRounding(displayId)
-            )
-            // TODO: b/362720497 - activating a desk with the intention to move a new task to
-            //  it means we may need to minimize something in the activating desk. Do so here
-            //  similar to how it's done in #bringDesktopAppsToFront instead of returning null.
-            null
-        } else {
-            bringDesktopAppsToFront(displayId, wct, newTask?.taskId)
+        if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
+            val taskIdToMinimize = bringDesktopAppsToFront(displayId, wct, newTask?.taskId)
+            return { transition ->
+                taskIdToMinimize?.let { minimizingTaskId ->
+                    addPendingMinimizeTransition(
+                        transition = transition,
+                        taskIdToMinimize = minimizingTaskId,
+                        minimizeReason = MinimizeReason.TASK_LIMIT,
+                    )
+                }
+            }
+        }
+        prepareForDeskActivation(displayId, wct)
+        desksOrganizer.activateDesk(wct, deskId)
+        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {
+            // TODO: 362720497 - do non-running tasks need to be restarted with |wct#startTask|?
+        }
+        taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(
+            doesAnyTaskRequireTaskbarRounding(displayId)
+        )
+        // TODO: b/362720497 - activating a desk with the intention to move a new task to
+        //  it means we may need to minimize something in the activating desk. Do so here
+        //  similar to how it's done in #bringDesktopAppsToFront.
+        return { transition ->
+            val activateDeskTransition =
+                if (newTaskIdInFront != null) {
+                    DeskTransition.ActiveDeskWithTask(
+                        token = transition,
+                        displayId = displayId,
+                        deskId = deskId,
+                        enterTaskId = newTaskIdInFront,
+                    )
+                } else {
+                    DeskTransition.ActivateDesk(
+                        token = transition,
+                        displayId = displayId,
+                        deskId = deskId,
+                    )
+                }
+            desksTransitionObserver.addPendingTransition(activateDeskTransition)
         }
     }
 
     /** Activates the given desk. */
     fun activateDesk(deskId: Int, remoteTransition: RemoteTransition? = null) {
-        val displayId = taskRepository.getDisplayForDesk(deskId)
         val wct = WindowContainerTransaction()
-        addDeskActivationChanges(deskId, wct)
+        val runOnTransitStart = addDeskActivationChanges(deskId, wct)
 
         val transitionType = transitionType(remoteTransition)
         val handler =
@@ -2814,15 +2804,7 @@ class DesktopTasksController(
 
         val transition = transitions.startTransition(transitionType, wct, handler)
         handler?.setTransition(transition)
-        if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-            desksTransitionObserver.addPendingTransition(
-                DeskTransition.ActivateDesk(
-                    token = transition,
-                    displayId = displayId,
-                    deskId = deskId,
-                )
-            )
-        }
+        runOnTransitStart?.invoke(transition)
 
         desktopModeEnterExitTransitionListener?.onEnterDesktopModeTransitionStarted(
             FREEFORM_ANIMATION_DURATION
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
index f9ab359e952d..b9b4d9e4bbd8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
@@ -268,6 +268,7 @@ class DesktopTasksLimiter(
         // If it's a running task, reorder it to back.
         taskIdToMinimize
             ?.let { shellTaskOrganizer.getRunningTaskInfo(it) }
+            // TODO: b/391485148 - this won't really work with multi-desks enabled.
             ?.let { wct.reorder(it.token, /* onTop= */ false) }
         return taskIdToMinimize
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index ea220c4135df..13a2e501f72e 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -2913,11 +2913,10 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         verify(desksOrganizer).activateDesk(any(), eq(targetDeskId))
         verify(desksTransitionsObserver)
             .addPendingTransition(
-                DeskTransition.ActiveDeskWithTask(
+                DeskTransition.ActivateDesk(
                     token = transition,
                     displayId = SECOND_DISPLAY,
                     deskId = targetDeskId,
-                    enterTaskId = task.taskId,
                 )
             )
     }