Clear all starting window while unregister organizer.

The starting window view might left in client's ShellTaskOrganizer,
clear all of the starting windows view while unregister task organizer
to preventing those views from left.
Consolidating the remove window methods together.

Bug: 207175662
Test: build/flash
Test: atest StartingSurfaceDrawerTests SplashscreenTests
Change-Id: I93e4862b7f866c7c73d9b3999a74e03784848058
Merged-In: I93e4862b7f866c7c73d9b3999a74e03784848058

4 files changed, 80 insertions, 15 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index 75bc46125324..26a2f5ab0723 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -200,6 +200,14 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
         }
     }
 
+    @Override
+    public void unregisterOrganizer() {
+        super.unregisterOrganizer();
+        if (mStartingWindow != null) {
+            mStartingWindow.clearAllWindows();
+        }
+    }
+
     public void createRootTask(int displayId, int windowingMode, TaskListener listener) {
         ProtoLog.v(WM_SHELL_TASK_ORG, "createRootTask() displayId=%d winMode=%d listener=%s",
                 displayId, windowingMode, listener.toString());
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
index bd4869670bec..270107c01335 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
@@ -134,7 +134,8 @@ public class StartingSurfaceDrawer {
         mDisplayManager.getDisplay(DEFAULT_DISPLAY);
     }
 
-    private final SparseArray<StartingWindowRecord> mStartingWindowRecords = new SparseArray<>();
+    @VisibleForTesting
+    final SparseArray<StartingWindowRecord> mStartingWindowRecords = new SparseArray<>();
 
     /**
      * Records of {@link SurfaceControlViewHost} where the splash screen icon animation is
@@ -464,8 +465,24 @@ public class StartingSurfaceDrawer {
             Slog.d(TAG, "Task start finish, remove starting surface for task "
                     + removalInfo.taskId);
         }
-        removeWindowSynced(removalInfo);
+        removeWindowSynced(removalInfo, false /* immediately */);
+    }
 
+    /**
+     * Clear all starting windows immediately.
+     */
+    public void clearAllWindows() {
+        if (DEBUG_SPLASH_SCREEN || DEBUG_TASK_SNAPSHOT) {
+            Slog.d(TAG, "Clear all starting windows immediately");
+        }
+        final int taskSize = mStartingWindowRecords.size();
+        final int[] taskIds = new int[taskSize];
+        for (int i = taskSize - 1; i >= 0; --i) {
+            taskIds[i] = mStartingWindowRecords.keyAt(i);
+        }
+        for (int i = taskSize - 1; i >= 0; --i) {
+            removeWindowNoAnimate(taskIds[i]);
+        }
     }
 
     /**
@@ -550,7 +567,8 @@ public class StartingSurfaceDrawer {
         return shouldSaveView;
     }
 
-    private void saveSplashScreenRecord(IBinder appToken, int taskId, View view,
+    @VisibleForTesting
+    void saveSplashScreenRecord(IBinder appToken, int taskId, View view,
             @StartingWindowType int suggestType) {
         final StartingWindowRecord tView = new StartingWindowRecord(appToken, view,
                 null/* TaskSnapshotWindow */, suggestType);
@@ -559,19 +577,18 @@ public class StartingSurfaceDrawer {
 
     private void removeWindowNoAnimate(int taskId) {
         mTmpRemovalInfo.taskId = taskId;
-        removeWindowSynced(mTmpRemovalInfo);
+        removeWindowSynced(mTmpRemovalInfo, true /* immediately */);
     }
 
     void onImeDrawnOnTask(int taskId) {
         final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
         if (record != null && record.mTaskSnapshotWindow != null
                 && record.mTaskSnapshotWindow.hasImeSurface()) {
-            record.mTaskSnapshotWindow.removeImmediately();
+            removeWindowNoAnimate(taskId);
         }
-        mStartingWindowRecords.remove(taskId);
     }
 
-    protected void removeWindowSynced(StartingWindowRemovalInfo removalInfo) {
+    protected void removeWindowSynced(StartingWindowRemovalInfo removalInfo, boolean immediately) {
         final int taskId = removalInfo.taskId;
         final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
         if (record != null) {
@@ -580,7 +597,8 @@ public class StartingSurfaceDrawer {
                     Slog.v(TAG, "Removing splash screen window for task: " + taskId);
                 }
                 if (record.mContentView != null) {
-                    if (record.mSuggestType == STARTING_WINDOW_TYPE_LEGACY_SPLASH_SCREEN) {
+                    if (immediately
+                            || record.mSuggestType == STARTING_WINDOW_TYPE_LEGACY_SPLASH_SCREEN) {
                         removeWindowInner(record.mDecorView, false);
                     } else {
                         if (removalInfo.playRevealAnimation) {
@@ -604,8 +622,12 @@ public class StartingSurfaceDrawer {
                 if (DEBUG_TASK_SNAPSHOT) {
                     Slog.v(TAG, "Removing task snapshot window for " + taskId);
                 }
-                record.mTaskSnapshotWindow.scheduleRemove(
-                        () -> mStartingWindowRecords.remove(taskId), removalInfo.deferRemoveForIme);
+                if (immediately) {
+                    record.mTaskSnapshotWindow.removeImmediately();
+                } else {
+                    record.mTaskSnapshotWindow.scheduleRemove(() ->
+                            mStartingWindowRecords.remove(taskId), removalInfo.deferRemoveForIme);
+                }
             }
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
index e98a3e87c0b7..b0a66059a466 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
@@ -29,9 +29,7 @@ import android.app.TaskInfo;
 import android.content.Context;
 import android.graphics.Color;
 import android.os.IBinder;
-import android.os.RemoteException;
 import android.os.Trace;
-import android.util.Slog;
 import android.util.SparseIntArray;
 import android.window.StartingWindowInfo;
 import android.window.StartingWindowInfo.StartingWindowType;
@@ -199,6 +197,18 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
     }
 
     /**
+     * Clear all starting window immediately, called this method when releasing the task organizer.
+     */
+    public void clearAllWindows() {
+        mSplashScreenExecutor.execute(() -> {
+            mStartingSurfaceDrawer.clearAllWindows();
+            synchronized (mTaskBackgroundColors) {
+                mTaskBackgroundColors.clear();
+            }
+        });
+    }
+
+    /**
      * The interface for calls from outside the Shell, within the host process.
      */
     private class StartingSurfaceImpl implements StartingSurface {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
index 70b7c6793492..d92b12e60780 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
@@ -31,6 +31,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -117,16 +118,19 @@ public class StartingSurfaceDrawerTests {
                 WindowManager.LayoutParams params, int suggestType) {
             // listen for addView
             mAddWindowForTask = taskId;
+            saveSplashScreenRecord(appToken, taskId, view, suggestType);
             // Do not wait for background color
             return false;
         }
 
         @Override
-        protected void removeWindowSynced(StartingWindowRemovalInfo removalInfo) {
+        protected void removeWindowSynced(StartingWindowRemovalInfo removalInfo,
+                boolean immediately) {
             // listen for removeView
             if (mAddWindowForTask == removalInfo.taskId) {
                 mAddWindowForTask = 0;
             }
+            mStartingWindowRecords.remove(removalInfo.taskId);
         }
     }
 
@@ -179,7 +183,7 @@ public class StartingSurfaceDrawerTests {
         removalInfo.taskId = windowInfo.taskInfo.taskId;
         mStartingSurfaceDrawer.removeStartingWindow(removalInfo);
         waitHandlerIdle(mTestHandler);
-        verify(mStartingSurfaceDrawer).removeWindowSynced(any());
+        verify(mStartingSurfaceDrawer).removeWindowSynced(any(), eq(false));
         assertEquals(mStartingSurfaceDrawer.mAddWindowForTask, 0);
     }
 
@@ -267,11 +271,32 @@ public class StartingSurfaceDrawerTests {
 
             // Verify the task snapshot with IME snapshot will be removed when received the real IME
             // drawn callback.
+            // makeTaskSnapshotWindow shall call removeWindowSynced before there add a new
+            // StartingWindowRecord for the task.
             mStartingSurfaceDrawer.onImeDrawnOnTask(1);
-            verify(mockSnapshotWindow).removeImmediately();
+            verify(mStartingSurfaceDrawer, times(2))
+                    .removeWindowSynced(any(), eq(true));
         }
     }
 
+    @Test
+    public void testClearAllWindows() {
+        final int taskId = 1;
+        final StartingWindowInfo windowInfo =
+                createWindowInfo(taskId, android.R.style.Theme);
+        mStartingSurfaceDrawer.addSplashScreenStartingWindow(windowInfo, mBinder,
+                STARTING_WINDOW_TYPE_SPLASH_SCREEN);
+        waitHandlerIdle(mTestHandler);
+        verify(mStartingSurfaceDrawer).addWindow(eq(taskId), eq(mBinder), any(), any(), any(),
+                eq(STARTING_WINDOW_TYPE_SPLASH_SCREEN));
+        assertEquals(mStartingSurfaceDrawer.mAddWindowForTask, taskId);
+
+        mStartingSurfaceDrawer.clearAllWindows();
+        waitHandlerIdle(mTestHandler);
+        verify(mStartingSurfaceDrawer).removeWindowSynced(any(), eq(true));
+        assertEquals(mStartingSurfaceDrawer.mStartingWindowRecords.size(), 0);
+    }
+
     private StartingWindowInfo createWindowInfo(int taskId, int themeResId) {
         StartingWindowInfo windowInfo = new StartingWindowInfo();
         final ActivityInfo info = new ActivityInfo();