Remove snapshot when screen gets rotated in mean time

We also need to wait for windows that are not gone for layout, as
otherwise we'd unfreeze before the surface is actually created.

Test: go/wm-smoke
Test: Open Camera in landscape but home screen locked to portrait
Change-Id: I0fc7a016445e18af1eead292665702b768b4e95b
Fixes: 38261533
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotSurface.java b/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
index 551e3bf..a96d224 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
@@ -125,6 +125,7 @@
     private final Paint mBackgroundPaint = new Paint();
     private final int mStatusBarColor;
     @VisibleForTesting final SystemBarBackgroundPainter mSystemBarBackgroundPainter;
+    private final int mOrientationOnCreation;
 
     static TaskSnapshotSurface create(WindowManagerService service, AppWindowToken token,
             TaskSnapshot snapshot) {
@@ -146,6 +147,7 @@
         final int sysUiVis;
         final int windowFlags;
         final int windowPrivateFlags;
+        final int currentOrientation;
         synchronized (service.mWindowMap) {
             final WindowState mainWindow = token.findMainWindow();
             if (mainWindow == null) {
@@ -183,6 +185,7 @@
             } else {
                 taskBounds = null;
             }
+            currentOrientation = mainWindow.getConfiguration().orientation;
         }
         try {
             final int res = session.addToDisplay(window, window.mSeq, layoutParams,
@@ -197,7 +200,8 @@
         }
         final TaskSnapshotSurface snapshotSurface = new TaskSnapshotSurface(service, window,
                 surface, snapshot, layoutParams.getTitle(), backgroundColor, statusBarColor,
-                navigationBarColor, sysUiVis, windowFlags, windowPrivateFlags, taskBounds);
+                navigationBarColor, sysUiVis, windowFlags, windowPrivateFlags, taskBounds,
+                currentOrientation);
         window.setOuter(snapshotSurface);
         try {
             session.relayout(window, window.mSeq, layoutParams, -1, -1, View.VISIBLE, 0, tmpFrame,
@@ -215,7 +219,7 @@
     TaskSnapshotSurface(WindowManagerService service, Window window, Surface surface,
             TaskSnapshot snapshot, CharSequence title, int backgroundColor, int statusBarColor,
             int navigationBarColor, int sysUiVis, int windowFlags, int windowPrivateFlags,
-            Rect taskBounds) {
+            Rect taskBounds, int currentOrientation) {
         mService = service;
         mHandler = new Handler(mService.mH.getLooper());
         mSession = WindowManagerGlobal.getWindowSession();
@@ -228,6 +232,7 @@
         mSystemBarBackgroundPainter = new SystemBarBackgroundPainter(windowFlags,
                 windowPrivateFlags, sysUiVis, statusBarColor, navigationBarColor);
         mStatusBarColor = statusBarColor;
+        mOrientationOnCreation = currentOrientation;
     }
 
     @Override
@@ -394,6 +399,7 @@
     static class Window extends BaseIWindow {
 
         private TaskSnapshotSurface mOuter;
+
         public void setOuter(TaskSnapshotSurface outer) {
             mOuter = outer;
         }
@@ -403,6 +409,15 @@
                 Rect stableInsets, Rect outsets, boolean reportDraw,
                 MergedConfiguration mergedConfiguration, Rect backDropFrame, boolean forceLayout,
                 boolean alwaysConsumeNavBar, int displayId) {
+            if (mergedConfiguration != null && mOuter != null
+                    && mOuter.mOrientationOnCreation
+                            != mergedConfiguration.getMergedConfiguration().orientation) {
+
+                // The orientation of the screen is changing. We better remove the snapshot ASAP as
+                // we are going to wait on the new window in any case to unfreeze the screen, and
+                // the starting window is not needed anymore.
+                sHandler.post(mOuter::remove);
+            }
             if (reportDraw) {
                 sHandler.obtainMessage(MSG_REPORT_DRAW, mOuter).sendToTarget();
             }