Remove starting window on the corresponding task organizer.

There can be multiple task organizer(s) register to core and exist at
the same time, to ensure the starting window can removed at the correct
client without mass up the starting window lifecycle, record the
organizer which create the starting window, so it can be removed on it
correctly.

Bug: 295621224
Test: presubmit
Change-Id: I207a25f6c8280a32e90b475902e2150929103192
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index d0f86c0..2794aab 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -2703,7 +2703,8 @@
 
     private void requestCopySplashScreen() {
         mTransferringSplashScreenState = TRANSFER_SPLASH_SCREEN_COPYING;
-        if (!mAtmService.mTaskOrganizerController.copySplashScreenView(getTask())) {
+        if (mStartingSurface == null || !mAtmService.mTaskOrganizerController.copySplashScreenView(
+                getTask(), mStartingSurface.mTaskOrganizer)) {
             mTransferringSplashScreenState = TRANSFER_SPLASH_SCREEN_FINISH;
             removeStartingWindow();
         }
@@ -2771,7 +2772,8 @@
                 && (mTransferringSplashScreenState == TRANSFER_SPLASH_SCREEN_FINISH
                 || mTransferringSplashScreenState == TRANSFER_SPLASH_SCREEN_IDLE)) {
             ProtoLog.v(WM_DEBUG_STARTING_WINDOW, "Cleaning splash screen token=%s", this);
-            mAtmService.mTaskOrganizerController.onAppSplashScreenViewRemoved(getTask());
+            mAtmService.mTaskOrganizerController.onAppSplashScreenViewRemoved(getTask(),
+                    mStartingSurface != null ? mStartingSurface.mTaskOrganizer : null);
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/StartingSurfaceController.java b/services/core/java/com/android/server/wm/StartingSurfaceController.java
index 0bb773a..a55c232 100644
--- a/services/core/java/com/android/server/wm/StartingSurfaceController.java
+++ b/services/core/java/com/android/server/wm/StartingSurfaceController.java
@@ -40,6 +40,7 @@
 import android.content.pm.ApplicationInfo;
 import android.os.UserHandle;
 import android.util.Slog;
+import android.window.ITaskOrganizer;
 import android.window.SplashScreenView;
 import android.window.TaskSnapshot;
 
@@ -79,12 +80,13 @@
     }
 
     StartingSurface createSplashScreenStartingSurface(ActivityRecord activity, int theme) {
-
         synchronized (mService.mGlobalLock) {
             final Task task = activity.getTask();
-            if (task != null && mService.mAtmService.mTaskOrganizerController.addStartingWindow(
-                    task, activity, theme, null /* taskSnapshot */)) {
-                return new StartingSurface(task);
+            final TaskOrganizerController controller =
+                    mService.mAtmService.mTaskOrganizerController;
+            if (task != null && controller.addStartingWindow(task, activity, theme,
+                    null /* taskSnapshot */)) {
+                return new StartingSurface(task, controller.getTaskOrganizer());
             }
         }
         return null;
@@ -166,9 +168,12 @@
                 activity.mDisplayContent.handleTopActivityLaunchingInDifferentOrientation(
                         activity, false /* checkOpening */);
             }
-                mService.mAtmService.mTaskOrganizerController.addStartingWindow(task,
-                        activity, 0 /* launchTheme */, taskSnapshot);
-            return new StartingSurface(task);
+            final TaskOrganizerController controller =
+                    mService.mAtmService.mTaskOrganizerController;
+            if (controller.addStartingWindow(task, activity, 0 /* launchTheme */, taskSnapshot)) {
+                return new StartingSurface(task, controller.getTaskOrganizer());
+            }
+            return null;
         }
     }
 
@@ -256,9 +261,12 @@
 
     final class StartingSurface {
         private final Task mTask;
+        // The task organizer which hold the client side reference of this surface.
+        final ITaskOrganizer mTaskOrganizer;
 
-        StartingSurface(Task task) {
+        StartingSurface(Task task, ITaskOrganizer taskOrganizer) {
             mTask = task;
+            mTaskOrganizer = taskOrganizer;
         }
 
         /**
@@ -268,7 +276,8 @@
          */
         public void remove(boolean animate) {
             synchronized (mService.mGlobalLock) {
-                mService.mAtmService.mTaskOrganizerController.removeStartingWindow(mTask, animate);
+                mService.mAtmService.mTaskOrganizerController.removeStartingWindow(mTask,
+                        mTaskOrganizer, animate);
             }
         }
     }
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index b72d027..baa101e 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -652,7 +652,7 @@
         if (rootTask == null || activity.mStartingData == null) {
             return false;
         }
-        final ITaskOrganizer lastOrganizer = mTaskOrganizers.peekLast();
+        final ITaskOrganizer lastOrganizer = getTaskOrganizer();
         if (lastOrganizer == null) {
             return false;
         }
@@ -672,12 +672,13 @@
         return true;
     }
 
-    void removeStartingWindow(Task task, boolean prepareAnimation) {
+    void removeStartingWindow(Task task, ITaskOrganizer taskOrganizer, boolean prepareAnimation) {
         final Task rootTask = task.getRootTask();
         if (rootTask == null) {
             return;
         }
-        final ITaskOrganizer lastOrganizer = mTaskOrganizers.peekLast();
+        final ITaskOrganizer lastOrganizer = taskOrganizer != null ? taskOrganizer
+                : getTaskOrganizer();
         if (lastOrganizer == null) {
             return;
         }
@@ -771,12 +772,13 @@
         }
     }
 
-    boolean copySplashScreenView(Task task) {
+    boolean copySplashScreenView(Task task, ITaskOrganizer taskOrganizer) {
         final Task rootTask = task.getRootTask();
         if (rootTask == null) {
             return false;
         }
-        final ITaskOrganizer lastOrganizer = mTaskOrganizers.peekLast();
+        final ITaskOrganizer lastOrganizer = taskOrganizer != null ? taskOrganizer
+                : getTaskOrganizer();
         if (lastOrganizer == null) {
             return false;
         }
@@ -799,12 +801,12 @@
      * @see com.android.wm.shell.ShellTaskOrganizer#onAppSplashScreenViewRemoved(int)
      * @see SplashScreenView#remove()
      */
-    public void onAppSplashScreenViewRemoved(Task task) {
+    public void onAppSplashScreenViewRemoved(Task task, ITaskOrganizer organizer) {
         final Task rootTask = task.getRootTask();
         if (rootTask == null) {
             return;
         }
-        final ITaskOrganizer lastOrganizer = mTaskOrganizers.peekLast();
+        final ITaskOrganizer lastOrganizer = organizer != null ? organizer : getTaskOrganizer();
         if (lastOrganizer == null) {
             return;
         }