diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 350a333..bd37934 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -394,6 +394,10 @@
                     startingWindow.mPolicyVisibility = false;
                     startingWindow.mPolicyVisibilityAfterAnim = false;
                 }
+
+                // We are becoming visible, so better freeze the screen with the windows that are
+                // getting visible so we also wait for them.
+                forAllWindows(mService::makeWindowFreezingScreenIfNeededLocked, true);
             }
 
             if (DEBUG_APP_TRANSITIONS) Slog.v(TAG_WM, "setVisibility: " + this
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();
             }
diff --git a/services/core/java/com/android/server/wm/WindowStateAnimator.java b/services/core/java/com/android/server/wm/WindowStateAnimator.java
index 73f8d27..33cb908 100644
--- a/services/core/java/com/android/server/wm/WindowStateAnimator.java
+++ b/services/core/java/com/android/server/wm/WindowStateAnimator.java
@@ -1524,7 +1524,10 @@
     void prepareSurfaceLocked(final boolean recoveringMemory) {
         final WindowState w = mWin;
         if (!hasSurface()) {
-            if (w.mOrientationChanging) {
+
+            // There is no need to wait for an animation change if our window is gone for layout
+            // already as we'll never be visible.
+            if (w.mOrientationChanging && w.isGoneForLayoutLw()) {
                 if (DEBUG_ORIENTATION) {
                     Slog.v(TAG, "Orientation change skips hidden " + w);
                 }
@@ -1557,13 +1560,11 @@
             hide("prepareSurfaceLocked");
             mWallpaperControllerLocked.hideWallpapers(w);
 
-            // If we are waiting for this window to handle an
-            // orientation change, well, it is hidden, so
-            // doesn't really matter.  Note that this does
-            // introduce a potential glitch if the window
-            // becomes unhidden before it has drawn for the
-            // new orientation.
-            if (w.mOrientationChanging) {
+            // If we are waiting for this window to handle an orientation change. If this window is
+            // really hidden (gone for layout), there is no point in still waiting for it.
+            // Note that this does introduce a potential glitch if the window becomes unhidden
+            // before it has drawn for the new orientation.
+            if (w.mOrientationChanging && w.isGoneForLayoutLw()) {
                 w.mOrientationChanging = false;
                 if (DEBUG_ORIENTATION) Slog.v(TAG,
                         "Orientation change skips hidden " + w);
@@ -1630,18 +1631,19 @@
             displayed = true;
         }
 
-        if (displayed) {
-            if (w.mOrientationChanging) {
-                if (!w.isDrawnLw()) {
-                    mAnimator.mBulkUpdateParams &= ~SET_ORIENTATION_CHANGE_COMPLETE;
-                    mAnimator.mLastWindowFreezeSource = w;
-                    if (DEBUG_ORIENTATION) Slog.v(TAG,
-                            "Orientation continue waiting for draw in " + w);
-                } else {
-                    w.mOrientationChanging = false;
-                    if (DEBUG_ORIENTATION) Slog.v(TAG, "Orientation change complete in " + w);
-                }
+        if (w.mOrientationChanging) {
+            if (!w.isDrawnLw()) {
+                mAnimator.mBulkUpdateParams &= ~SET_ORIENTATION_CHANGE_COMPLETE;
+                mAnimator.mLastWindowFreezeSource = w;
+                if (DEBUG_ORIENTATION) Slog.v(TAG,
+                        "Orientation continue waiting for draw in " + w);
+            } else {
+                w.mOrientationChanging = false;
+                if (DEBUG_ORIENTATION) Slog.v(TAG, "Orientation change complete in " + w);
             }
+        }
+
+        if (displayed) {
             w.mToken.hasVisible = true;
         }
     }
diff --git a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
index e2868d7..4288eac 100644
--- a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
+++ b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
@@ -62,7 +62,8 @@
         final TaskSnapshot snapshot = new TaskSnapshot(buffer,
                 ORIENTATION_PORTRAIT, contentInsets, false, 1.0f);
         mSurface = new TaskSnapshotSurface(sWm, new Window(), new Surface(), snapshot, "Test",
-                Color.WHITE, Color.RED, Color.BLUE, sysuiVis, windowFlags, 0, taskBounds);
+                Color.WHITE, Color.RED, Color.BLUE, sysuiVis, windowFlags, 0, taskBounds,
+                ORIENTATION_PORTRAIT);
     }
 
     private void setupSurface(int width, int height) {