Introduce display specific all window drawn waiting (3/N)

This change modified the waitForAllWindowsDrawn API to wait for
windows on a specific display. INVALID_DISPLAY indicates to wait on all
visible windows on all displays.

To achieve that, it introduced a callback HashMap in WMS and a window
container as a parameter in the waiting related messages. The actual
window registration is done by WindowContainere

Bug: 141528152
Bug: 122726344
Bug: 134531136
Test: go/wm-smoke
Change-Id: Ie6c8acfccef6dcad34e208acc37dd87b1e18b0e1
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 88b1793..4d3f853 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -4572,7 +4572,7 @@
         // ... eventually calls finishWindowsDrawn which will finalize our screen turn on
         // as well as enabling the orientation change logic/sensor.
         mWindowManagerInternal.waitForAllWindowsDrawn(mWindowManagerDrawCallback,
-                WAITING_FOR_DRAWN_TIMEOUT);
+                WAITING_FOR_DRAWN_TIMEOUT, INVALID_DISPLAY);
     }
 
     // Called on the DisplayManager's DisplayPowerController thread.
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 6d9a008..c97bd14 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -132,7 +132,6 @@
 import static com.android.server.wm.WindowState.EXCLUSION_LEFT;
 import static com.android.server.wm.WindowState.EXCLUSION_RIGHT;
 import static com.android.server.wm.WindowState.RESIZE_HANDLE_WIDTH_IN_DP;
-import static com.android.server.wm.WindowStateAnimator.DRAW_PENDING;
 import static com.android.server.wm.WindowStateAnimator.READY_TO_SHOW;
 import static com.android.server.wm.utils.RegionUtils.forEachRectReverse;
 import static com.android.server.wm.utils.RegionUtils.rectListToRegion;
@@ -3514,19 +3513,6 @@
         mWmService.mWindowPlacerLocked.performSurfacePlacement();
     }
 
-    void waitForAllWindowsDrawn() {
-        final WindowManagerPolicy policy = mWmService.mPolicy;
-        forAllWindows(w -> {
-            final boolean keyguard = policy.isKeyguardHostWindow(w.mAttrs);
-            if (w.isVisibleLw() && (w.mAppToken != null || keyguard)) {
-                w.mWinAnimator.mDrawState = DRAW_PENDING;
-                // Force add to mResizingWindows.
-                w.resetLastContentInsets();
-                mWmService.mWaitingForDrawn.add(w);
-            }
-        }, true /* traverseTopToBottom */);
-    }
-
     // TODO: Super crazy long method that should be broken down...
     void applySurfaceChangesTransaction(boolean recoveringMemory) {
         final WindowSurfacePlacer surfacePlacer = mWmService.mWindowPlacerLocked;
diff --git a/services/core/java/com/android/server/wm/DockedStackDividerController.java b/services/core/java/com/android/server/wm/DockedStackDividerController.java
index ae3b5f2..d3e42901 100644
--- a/services/core/java/com/android/server/wm/DockedStackDividerController.java
+++ b/services/core/java/com/android/server/wm/DockedStackDividerController.java
@@ -819,9 +819,12 @@
 
         // We put all tasks into drag resizing mode - wait until all of them have completed the
         // drag resizing switch.
-        if (!mService.mWaitingForDrawn.isEmpty()) {
-            mService.mH.removeMessages(H.WAITING_FOR_DRAWN_TIMEOUT);
-            mService.mH.sendEmptyMessageDelayed(H.WAITING_FOR_DRAWN_TIMEOUT,
+        final Runnable existingWaitingForDrwanCallback =
+                mService.mWaitingForDrawnCallbacks.get(mService.mRoot);
+        if (existingWaitingForDrwanCallback != null) {
+            mService.mH.removeMessages(H.WAITING_FOR_DRAWN_TIMEOUT, mService.mRoot);
+            mService.mH.sendMessageDelayed(mService.mH.obtainMessage(H.WAITING_FOR_DRAWN_TIMEOUT,
+                    mService.mRoot),
                     IME_ADJUST_DRAWN_TIMEOUT);
             mAnimationStartDelayed = true;
             if (imeWin != null) {
@@ -838,10 +841,8 @@
             // still gets executed.
             // TODO: Have a real system where we can wait on different windows to be drawn with
             // different callbacks.
-            if (mService.mWaitingForDrawnCallback != null) {
-                mService.mWaitingForDrawnCallback.run();
-            }
-            mService.mWaitingForDrawnCallback = () -> {
+            existingWaitingForDrwanCallback.run();
+            mService.mWaitingForDrawnCallbacks.put(mService.mRoot, () -> {
                 synchronized (mService.mGlobalLock) {
                     mAnimationStartDelayed = false;
                     if (mDelayedImeWin != null) {
@@ -863,7 +864,7 @@
                     notifyAdjustedForImeChanged(
                             mAdjustedForIme || mAdjustedForDivider, duration);
                 }
-            };
+            });
         } else {
             notifyAdjustedForImeChanged(
                     adjustedForIme || adjustedForDivider, IME_ADJUST_ANIM_DURATION);
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 78fcb37..b775002 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -720,7 +720,7 @@
             mUpdateRotation = updateRotationUnchecked();
         }
 
-        if (mWmService.mWaitingForDrawnCallback != null
+        if (!mWmService.mWaitingForDrawnCallbacks.isEmpty()
                 || (mOrientationChangeComplete && !isLayoutNeeded()
                 && !mUpdateRotation)) {
             mWmService.checkDrawnWindowsLocked();
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 264efe0..6de30a9 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -28,6 +28,7 @@
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_ANIM;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
+import static com.android.server.wm.WindowStateAnimator.DRAW_PENDING;
 
 import android.annotation.CallSuper;
 import android.annotation.IntDef;
@@ -47,9 +48,11 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ToBooleanFunction;
+import com.android.server.policy.WindowManagerPolicy;
 import com.android.server.wm.SurfaceAnimator.Animatable;
 
 import java.io.PrintWriter;
+import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.LinkedList;
 import java.util.function.Consumer;
@@ -117,6 +120,11 @@
     private final Transaction mPendingTransaction;
 
     /**
+     * Windows that clients are waiting to have drawn.
+     */
+    final ArrayList<WindowState> mWaitingForDrawn = new ArrayList<>();
+
+    /**
      * Applied as part of the animation pass in "prepareSurfaces".
      */
     protected final SurfaceAnimator mSurfaceAnimator;
@@ -1392,6 +1400,19 @@
         }
     }
 
+    void waitForAllWindowsDrawn() {
+        final WindowManagerPolicy policy = mWmService.mPolicy;
+        forAllWindows(w -> {
+            final boolean keyguard = policy.isKeyguardHostWindow(w.mAttrs);
+            if (w.isVisibleLw() && (w.mAppToken != null || keyguard)) {
+                w.mWinAnimator.mDrawState = DRAW_PENDING;
+                // Force add to mResizingWindows.
+                w.resetLastContentInsets();
+                mWaitingForDrawn.add(w);
+            }
+        }, true /* traverseTopToBottom */);
+    }
+
     Dimmer getDimmer() {
         if (mParent == null) {
             return null;
diff --git a/services/core/java/com/android/server/wm/WindowManagerInternal.java b/services/core/java/com/android/server/wm/WindowManagerInternal.java
index abbd1ab..273d9a4 100644
--- a/services/core/java/com/android/server/wm/WindowManagerInternal.java
+++ b/services/core/java/com/android/server/wm/WindowManagerInternal.java
@@ -313,10 +313,15 @@
     public abstract void showGlobalActions();
 
     /**
-     * Invalidate all visible windows. Then report back on the callback once all windows have
-     * redrawn.
+     * Invalidate all visible windows on a given display, and report back on the callback when all
+     * windows have redrawn.
+     *
+     * @param callback reporting callback to be called when all windows have redrawn.
+     * @param timeout calls the callback anyway after the timeout.
+     * @param displayId waits for the windows on the given display, INVALID_DISPLAY to wait for all
+     *                  windows on all displays.
      */
-    public abstract void waitForAllWindowsDrawn(Runnable callback, long timeout);
+    public abstract void waitForAllWindowsDrawn(Runnable callback, long timeout, int displayId);
 
     /**
      * Overrides the display size.
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 1f45cfb..b26aac7 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -290,6 +290,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
@@ -568,13 +569,10 @@
     final ArrayList<WindowState> mForceRemoves = new ArrayList<>();
 
     /**
-     * Windows that clients are waiting to have drawn.
+     * The callbacks to make when the windows all have been drawn for a given
+     * {@link WindowContainer}.
      */
-    ArrayList<WindowState> mWaitingForDrawn = new ArrayList<>();
-    /**
-     * And the callback to make when they've all been drawn.
-     */
-    Runnable mWaitingForDrawnCallback;
+    final HashMap<WindowContainer, Runnable> mWaitingForDrawnCallbacks = new HashMap<>();
 
     /** List of window currently causing non-system overlay windows to be hidden. */
     private ArrayList<WindowState> mHidingNonSystemOverlayWindows = new ArrayList<>();
@@ -4736,12 +4734,12 @@
 
                 case WAITING_FOR_DRAWN_TIMEOUT: {
                     Runnable callback = null;
+                    final WindowContainer container = (WindowContainer) msg.obj;
                     synchronized (mGlobalLock) {
                         ProtoLog.w(WM_ERROR, "Timeout waiting for drawn: undrawn=%s",
-                                mWaitingForDrawn);
-                        mWaitingForDrawn.clear();
-                        callback = mWaitingForDrawnCallback;
-                        mWaitingForDrawnCallback = null;
+                                container.mWaitingForDrawn);
+                        container.mWaitingForDrawn.clear();
+                        callback = mWaitingForDrawnCallbacks.remove(container);
                     }
                     if (callback != null) {
                         callback.run();
@@ -4773,9 +4771,9 @@
                 }
                 case ALL_WINDOWS_DRAWN: {
                     Runnable callback;
+                    final WindowContainer container = (WindowContainer) msg.obj;
                     synchronized (mGlobalLock) {
-                        callback = mWaitingForDrawnCallback;
-                        mWaitingForDrawnCallback = null;
+                        callback = mWaitingForDrawnCallbacks.remove(container);
                     }
                     if (callback != null) {
                         callback.run();
@@ -5265,30 +5263,32 @@
     }
 
     void checkDrawnWindowsLocked() {
-        if (mWaitingForDrawn.isEmpty() || mWaitingForDrawnCallback == null) {
+        if (mWaitingForDrawnCallbacks.isEmpty()) {
             return;
         }
-        for (int j = mWaitingForDrawn.size() - 1; j >= 0; j--) {
-            WindowState win = mWaitingForDrawn.get(j);
-            ProtoLog.i(WM_DEBUG_SCREEN_ON,
+        mWaitingForDrawnCallbacks.forEach((container, callback) -> {
+            for (int j = container.mWaitingForDrawn.size() - 1; j >= 0; j--) {
+                final WindowState win = (WindowState) container.mWaitingForDrawn.get(j);
+                ProtoLog.i(WM_DEBUG_SCREEN_ON,
                         "Waiting for drawn %s: removed=%b visible=%b mHasSurface=%b drawState=%d",
                         win, win.mRemoved, win.isVisibleLw(), win.mHasSurface,
                         win.mWinAnimator.mDrawState);
-            if (win.mRemoved || !win.mHasSurface || !win.isVisibleByPolicy()) {
-                // Window has been removed or hidden; no draw will now happen, so stop waiting.
-                ProtoLog.w(WM_DEBUG_SCREEN_ON, "Aborted waiting for drawn: %s", win);
-                mWaitingForDrawn.remove(win);
-            } else if (win.hasDrawnLw()) {
-                // Window is now drawn (and shown).
-                ProtoLog.d(WM_DEBUG_SCREEN_ON, "Window drawn win=%s", win);
-                mWaitingForDrawn.remove(win);
+                if (win.mRemoved || !win.mHasSurface || !win.isVisibleByPolicy()) {
+                    // Window has been removed or hidden; no draw will now happen, so stop waiting.
+                    ProtoLog.w(WM_DEBUG_SCREEN_ON, "Aborted waiting for drawn: %s", win);
+                    container.mWaitingForDrawn.remove(win);
+                } else if (win.hasDrawnLw()) {
+                    // Window is now drawn (and shown).
+                    ProtoLog.d(WM_DEBUG_SCREEN_ON, "Window drawn win=%s", win);
+                    container.mWaitingForDrawn.remove(win);
+                }
             }
-        }
-        if (mWaitingForDrawn.isEmpty()) {
-            ProtoLog.d(WM_DEBUG_SCREEN_ON, "All windows drawn!");
-            mH.removeMessages(H.WAITING_FOR_DRAWN_TIMEOUT);
-            mH.sendEmptyMessage(H.ALL_WINDOWS_DRAWN);
-        }
+            if (container.mWaitingForDrawn.isEmpty()) {
+                ProtoLog.d(WM_DEBUG_SCREEN_ON, "All windows drawn!");
+                mH.removeMessages(H.WAITING_FOR_DRAWN_TIMEOUT, container);
+                mH.sendMessage(mH.obtainMessage(H.ALL_WINDOWS_DRAWN, container));
+            }
+        });
     }
 
     void setHoldScreenLocked(final Session newHoldScreen) {
@@ -5934,13 +5934,18 @@
                 }
             }
         }
-        if (mWaitingForDrawn.size() > 0) {
+        if (!mWaitingForDrawnCallbacks.isEmpty()) {
             pw.println();
             pw.println("  Clients waiting for these windows to be drawn:");
-            for (int i=mWaitingForDrawn.size()-1; i>=0; i--) {
-                WindowState win = mWaitingForDrawn.get(i);
-                pw.print("  Waiting #"); pw.print(i); pw.print(' '); pw.print(win);
-            }
+            mWaitingForDrawnCallbacks.forEach((wc, callback) -> {
+                pw.print("  WindowContainer ");
+                pw.println(wc.getName());
+                for (int i = wc.mWaitingForDrawn.size() - 1; i >= 0; i--) {
+                    final WindowState win = (WindowState) wc.mWaitingForDrawn.get(i);
+                    pw.print("  Waiting #"); pw.print(i); pw.print(' '); pw.print(win);
+                }
+            });
+
         }
         pw.println();
         pw.print("  mGlobalConfiguration="); pw.println(mRoot.getConfiguration());
@@ -7096,17 +7101,25 @@
         }
 
         @Override
-        public void waitForAllWindowsDrawn(Runnable callback, long timeout) {
+        public void waitForAllWindowsDrawn(Runnable callback, long timeout, int displayId) {
+            final WindowContainer container = displayId == INVALID_DISPLAY
+                    ? mRoot : mRoot.getDisplayContent(displayId);
+            if (container == null) {
+                // The waiting container doesn't exist, no need to wait to run the callback. Run and
+                // return;
+                callback.run();
+                return;
+            }
             boolean allWindowsDrawn = false;
             synchronized (mGlobalLock) {
-                mWaitingForDrawnCallback = callback;
-                getDefaultDisplayContentLocked().waitForAllWindowsDrawn();
+                container.waitForAllWindowsDrawn();
                 mWindowPlacerLocked.requestTraversal();
-                mH.removeMessages(H.WAITING_FOR_DRAWN_TIMEOUT);
-                if (mWaitingForDrawn.isEmpty()) {
+                mH.removeMessages(H.WAITING_FOR_DRAWN_TIMEOUT, container);
+                if (container.mWaitingForDrawn.isEmpty()) {
                     allWindowsDrawn = true;
                 } else {
-                    mH.sendEmptyMessageDelayed(H.WAITING_FOR_DRAWN_TIMEOUT, timeout);
+                    mWaitingForDrawnCallbacks.put(container, callback);
+                    mH.sendNewMessageDelayed(H.WAITING_FOR_DRAWN_TIMEOUT, container, timeout);
                     checkDrawnWindowsLocked();
                 }
             }
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 569d248..4d36127 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -3497,7 +3497,7 @@
     @Override
     void setWaitingForDrawnIfResizingChanged() {
         if (isDragResizeChanged()) {
-            mWmService.mWaitingForDrawn.add(this);
+            mWmService.mRoot.mWaitingForDrawn.add(this);
         }
         super.setWaitingForDrawnIfResizingChanged();
     }