Update task orgs with existing tasks

- Move stack task org calls up to task
- If the task organizer dies, reset the last cached wm mode and update
  all existing tasks when the task organizer is next re-registered. This
  only applies to the first task organizer for a given windowing mode
  to ensure that the testing strategy of only reporting new tasks to
  current task organizers still applies.

Bug: 150874492
Test: atest WmTests:TaskOrganizerTests
Test: Move pip, kill sysui, ensure sysui still controls pip task
Change-Id: Ib266f67243ea215940747d57022d3aa91d55dff8

packages/SystemUI/src/com/android/systemui/pip/phone/PipManager.java

@@ -232,6 +232,12 @@ public class PipManager implements BasePipManager, PipTaskOrganizer.PipTransitio
                 mTouchHandler.getMotionHelper());
         displayController.addDisplayChangingController(mRotationController);
 
+        // Ensure that we have the display info in case we get calls to update the bounds before the
+        // listener calls back
+        final DisplayInfo displayInfo = new DisplayInfo();
+        context.getDisplay().getDisplayInfo(displayInfo);
+        mPipBoundsHandler.onDisplayInfoChanged(displayInfo);
+
         try {
             ActivityTaskManager.getTaskOrganizerController().registerTaskOrganizer(
                     mPipTaskOrganizer, WINDOWING_MODE_PINNED);

services/core/java/com/android/server/wm/ActivityStack.java

@@ -153,7 +153,6 @@ import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
 import android.view.Display;
 import android.view.DisplayInfo;
-import android.view.ITaskOrganizer;
 import android.view.SurfaceControl;
 
 import com.android.internal.annotations.GuardedBy;
@@ -307,8 +306,6 @@ class ActivityStack extends Task {
     // TODO(task-hierarchy): remove when tiles can be actual parents
     TaskTile mTile = null;
 
-    private int mLastTaskOrganizerWindowingMode = -1;
-
     private final Handler mHandler;
 
     private class ActivityStackHandler extends Handler {
@@ -635,8 +632,6 @@ class ActivityStack extends Task {
 
         super.onConfigurationChanged(newParentConfig);
 
-        updateTaskOrganizerState();
-
         // Only need to update surface size here since the super method will handle updating
         // surface position.
         updateSurfaceSize(getPendingTransaction());
@@ -692,30 +687,6 @@ class ActivityStack extends Task {
         }
     }
 
-    void updateTaskOrganizerState() {
-        if (!isRootTask()) {
-            return;
-        }
-
-        final int windowingMode = getWindowingMode();
-        if (windowingMode == mLastTaskOrganizerWindowingMode) {
-            // If our windowing mode hasn't actually changed, then just stick
-            // with our old organizer. This lets us implement the semantic
-            // where SysUI can continue to manage it's old tasks
-            // while CTS temporarily takes over the registration.
-            return;
-        }
-        /*
-         * Different windowing modes may be managed by different task organizers. If
-         * getTaskOrganizer returns null, we still call setTaskOrganizer to
-         * make sure we clear it.
-         */
-        final ITaskOrganizer org =
-            mWmService.mAtmService.mTaskOrganizerController.getTaskOrganizer(windowingMode);
-        setTaskOrganizer(org);
-        mLastTaskOrganizerWindowingMode = windowingMode;
-    }
-
     @Override
     public void setWindowingMode(int windowingMode) {
         // Calling Task#setWindowingMode() for leaf task since this is the a specialization of

services/core/java/com/android/server/wm/Task.java

@@ -474,6 +474,7 @@ class Task extends WindowContainer<WindowContainer> {
      * taskAppeared callback, and emit a taskRemoved callback when the Task is vanished.
      */
     ITaskOrganizer mTaskOrganizer;
+    private int mLastTaskOrganizerWindowingMode = -1;
 
     /**
      * Last Picture-in-Picture params applicable to the task. Updated when the app
@@ -1929,6 +1930,7 @@ class Task extends WindowContainer<WindowContainer> {
         // TODO: Should also take care of Pip mode changes here.
 
         saveLaunchingStateIfNeeded();
+        updateTaskOrganizerState(false /* forceUpdate */);
     }
 
     /**
@@ -4018,6 +4020,39 @@ class Task extends WindowContainer<WindowContainer> {
     // Called on Binder death.
     void taskOrganizerDied() {
         mTaskOrganizer = null;
+        mLastTaskOrganizerWindowingMode = -1;
+    }
+
+    /**
+     * Called when the task state changes (ie. from windowing mode change) an the task organizer
+     * state should also be updated.
+     *
+     * @param forceUpdate Updates the task organizer to the one currently specified in the task
+     *                    org controller for the task's windowing mode, ignoring the cached
+     *                    windowing mode checks.
+     */
+    void updateTaskOrganizerState(boolean forceUpdate) {
+        if (!isRootTask()) {
+            return;
+        }
+
+        final int windowingMode = getWindowingMode();
+        if (!forceUpdate && windowingMode == mLastTaskOrganizerWindowingMode) {
+            // If our windowing mode hasn't actually changed, then just stick
+            // with our old organizer. This lets us implement the semantic
+            // where SysUI can continue to manage it's old tasks
+            // while CTS temporarily takes over the registration.
+            return;
+        }
+        /*
+         * Different windowing modes may be managed by different task organizers. If
+         * getTaskOrganizer returns null, we still call setTaskOrganizer to
+         * make sure we clear it.
+         */
+        final ITaskOrganizer org =
+                mWmService.mAtmService.mTaskOrganizerController.getTaskOrganizer(windowingMode);
+        setTaskOrganizer(org);
+        mLastTaskOrganizerWindowingMode = windowingMode;
     }
 
     @Override

services/core/java/com/android/server/wm/TaskOrganizerController.java

@@ -105,7 +105,7 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub
 
 
         TaskOrganizerState(ITaskOrganizer organizer, int windowingMode,
-                TaskOrganizerState replacing) {
+                @Nullable TaskOrganizerState replacing) {
             mOrganizer = organizer;
             mDeathRecipient = new DeathRecipient(organizer, windowingMode);
             try {
@@ -203,10 +203,27 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub
         final long origId = Binder.clearCallingIdentity();
         try {
             synchronized (mGlobalLock) {
+                if (getTaskOrganizer(windowingMode) != null) {
+                    Slog.w(TAG, "Task organizer already exists for windowing mode: "
+                            + windowingMode);
+                }
+                final TaskOrganizerState previousState =
+                        mTaskOrganizersForWindowingMode.get(windowingMode);
                 final TaskOrganizerState state = new TaskOrganizerState(organizer, windowingMode,
-                        mTaskOrganizersForWindowingMode.get(windowingMode));
+                        previousState);
                 mTaskOrganizersForWindowingMode.put(windowingMode, state);
                 mTaskOrganizerStates.put(organizer.asBinder(), state);
+
+                if (previousState == null) {
+                    // Only in the case where this is the root task organizer for the given
+                    // windowing mode, we add report all existing tasks in that mode to the new
+                    // task organizer.
+                    mService.mRootWindowContainer.forAllTasks((task) -> {
+                        if (task.getWindowingMode() == windowingMode) {
+                            task.updateTaskOrganizerState(true /* forceUpdate */);
+                        }
+                    });
+                }
             }
         } finally {
             Binder.restoreCallingIdentity(origId);

services/tests/wmtests/src/com/android/server/wm/TaskOrganizerTests.java

@@ -212,6 +212,16 @@ public class TaskOrganizerTests extends WindowTestsBase {
         verify(organizer, times(1)).taskVanished(any());
     }
 
+    @Test
+    public void testRegisterTaskOrganizerWithExistingTasks() throws RemoteException {
+        final ActivityStack stack = createTaskStackOnDisplay(mDisplayContent);
+        final Task task = createTaskInStack(stack, 0 /* userId */);
+        stack.setWindowingMode(WINDOWING_MODE_PINNED);
+
+        final ITaskOrganizer organizer = registerMockOrganizer(WINDOWING_MODE_PINNED);
+        verify(organizer, times(1)).taskAppeared(any());
+    }
+
     @Test
     public void testTaskTransaction() {
         removeGlobalMinSizeRestriction();