diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index cbded89..a88e2b0 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -702,8 +702,11 @@
                     finishKeyguardDrawn();
                     break;
                 case MSG_WINDOW_MANAGER_DRAWN_COMPLETE:
-                    if (DEBUG_WAKEUP) Slog.w(TAG, "Setting mWindowManagerDrawComplete");
-                    finishWindowsDrawn(msg.arg1);
+                    final int displayId = msg.arg1;
+                    if (DEBUG_WAKEUP) Slog.w(TAG, "All windows drawn on display " + displayId);
+                    Trace.asyncTraceEnd(Trace.TRACE_TAG_WINDOW_MANAGER,
+                            TRACE_WAIT_FOR_ALL_WINDOWS_DRAWN_METHOD, displayId /* cookie */);
+                    finishWindowsDrawn(displayId);
                     break;
                 case MSG_HIDE_BOOT_MESSAGE:
                     handleHideBootMessage();
@@ -4996,15 +4999,10 @@
         // ... eventually calls finishWindowsDrawn which will finalize our screen turn on
         // as well as enabling the orientation change logic/sensor.
         Trace.asyncTraceBegin(Trace.TRACE_TAG_WINDOW_MANAGER,
-                TRACE_WAIT_FOR_ALL_WINDOWS_DRAWN_METHOD, /* cookie= */ 0);
-        mWindowManagerInternal.waitForAllWindowsDrawn(() -> {
-            if (DEBUG_WAKEUP) Slog.i(TAG, "All windows ready for every display");
-            mHandler.sendMessage(mHandler.obtainMessage(MSG_WINDOW_MANAGER_DRAWN_COMPLETE,
-                    INVALID_DISPLAY, 0));
-
-            Trace.asyncTraceEnd(Trace.TRACE_TAG_WINDOW_MANAGER,
-                    TRACE_WAIT_FOR_ALL_WINDOWS_DRAWN_METHOD, /* cookie= */ 0);
-            }, WAITING_FOR_DRAWN_TIMEOUT, INVALID_DISPLAY);
+                TRACE_WAIT_FOR_ALL_WINDOWS_DRAWN_METHOD, INVALID_DISPLAY /* cookie */);
+        mWindowManagerInternal.waitForAllWindowsDrawn(mHandler.obtainMessage(
+                MSG_WINDOW_MANAGER_DRAWN_COMPLETE, INVALID_DISPLAY, 0),
+                WAITING_FOR_DRAWN_TIMEOUT, INVALID_DISPLAY);
     }
 
     // Called on the DisplayManager's DisplayPowerController thread.
@@ -5084,15 +5082,10 @@
             mScreenOnListeners.put(displayId, screenOnListener);
 
             Trace.asyncTraceBegin(Trace.TRACE_TAG_WINDOW_MANAGER,
-                    TRACE_WAIT_FOR_ALL_WINDOWS_DRAWN_METHOD, /* cookie= */ 0);
-            mWindowManagerInternal.waitForAllWindowsDrawn(() -> {
-                if (DEBUG_WAKEUP) Slog.i(TAG, "All windows ready for display: " + displayId);
-                mHandler.sendMessage(mHandler.obtainMessage(MSG_WINDOW_MANAGER_DRAWN_COMPLETE,
-                        displayId, 0));
-
-                Trace.asyncTraceEnd(Trace.TRACE_TAG_WINDOW_MANAGER,
-                        TRACE_WAIT_FOR_ALL_WINDOWS_DRAWN_METHOD, /* cookie= */ 0);
-            }, WAITING_FOR_DRAWN_TIMEOUT, displayId);
+                    TRACE_WAIT_FOR_ALL_WINDOWS_DRAWN_METHOD, displayId /* cookie */);
+            mWindowManagerInternal.waitForAllWindowsDrawn(mHandler.obtainMessage(
+                    MSG_WINDOW_MANAGER_DRAWN_COMPLETE, displayId, 0),
+                    WAITING_FOR_DRAWN_TIMEOUT, displayId);
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/WindowManagerInternal.java b/services/core/java/com/android/server/wm/WindowManagerInternal.java
index 9e7df00..805e7ff 100644
--- a/services/core/java/com/android/server/wm/WindowManagerInternal.java
+++ b/services/core/java/com/android/server/wm/WindowManagerInternal.java
@@ -30,6 +30,7 @@
 import android.hardware.display.DisplayManagerInternal;
 import android.os.Bundle;
 import android.os.IBinder;
+import android.os.Message;
 import android.util.Pair;
 import android.view.ContentRecordingSession;
 import android.view.Display;
@@ -515,12 +516,13 @@
      * 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 message The message will be sent when all windows have redrawn. Note that the message
+     *                must be obtained from handler, otherwise it will throw NPE.
      * @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, int displayId);
+    public abstract void waitForAllWindowsDrawn(Message message, 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 67572bf..d9ff91f 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -601,7 +601,7 @@
      * The callbacks to make when the windows all have been drawn for a given
      * {@link WindowContainer}.
      */
-    final HashMap<WindowContainer, Runnable> mWaitingForDrawnCallbacks = new HashMap<>();
+    final ArrayMap<WindowContainer<?>, Message> mWaitingForDrawnCallbacks = new ArrayMap<>();
 
     /** List of window currently causing non-system overlay windows to be hidden. */
     private ArrayList<WindowState> mHidingNonSystemOverlayWindows = new ArrayList<>();
@@ -5368,8 +5368,6 @@
         public static final int CLIENT_FREEZE_TIMEOUT = 30;
         public static final int NOTIFY_ACTIVITY_DRAWN = 32;
 
-        public static final int ALL_WINDOWS_DRAWN = 33;
-
         public static final int NEW_ANIMATOR_SCALE = 34;
 
         public static final int SHOW_EMULATOR_DISPLAY_OVERLAY = 36;
@@ -5491,7 +5489,7 @@
                 }
 
                 case WAITING_FOR_DRAWN_TIMEOUT: {
-                    Runnable callback = null;
+                    final Message callback;
                     final WindowContainer<?> container = (WindowContainer<?>) msg.obj;
                     synchronized (mGlobalLock) {
                         ProtoLog.w(WM_ERROR, "Timeout waiting for drawn: undrawn=%s",
@@ -5505,7 +5503,7 @@
                         callback = mWaitingForDrawnCallbacks.remove(container);
                     }
                     if (callback != null) {
-                        callback.run();
+                        callback.sendToTarget();
                     }
                     break;
                 }
@@ -5529,17 +5527,6 @@
                     }
                     break;
                 }
-                case ALL_WINDOWS_DRAWN: {
-                    Runnable callback;
-                    final WindowContainer container = (WindowContainer) msg.obj;
-                    synchronized (mGlobalLock) {
-                        callback = mWaitingForDrawnCallbacks.remove(container);
-                    }
-                    if (callback != null) {
-                        callback.run();
-                    }
-                    break;
-                }
                 case NEW_ANIMATOR_SCALE: {
                     float scale = getCurrentAnimatorScale();
                     ValueAnimator.setDurationScale(scale);
@@ -6097,7 +6084,8 @@
         if (mWaitingForDrawnCallbacks.isEmpty()) {
             return;
         }
-        mWaitingForDrawnCallbacks.forEach((container, callback) -> {
+        for (int i = mWaitingForDrawnCallbacks.size() - 1; i >= 0; i--) {
+            final WindowContainer<?> container = mWaitingForDrawnCallbacks.keyAt(i);
             for (int j = container.mWaitingForDrawn.size() - 1; j >= 0; j--) {
                 final WindowState win = (WindowState) container.mWaitingForDrawn.get(j);
                 ProtoLog.i(WM_DEBUG_SCREEN_ON,
@@ -6123,9 +6111,9 @@
             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));
+                mWaitingForDrawnCallbacks.removeAt(i).sendToTarget();
             }
-        });
+        }
     }
 
     private void traceStartWaitingForWindowDrawn(WindowState window) {
@@ -7811,13 +7799,14 @@
         }
 
         @Override
-        public void waitForAllWindowsDrawn(Runnable callback, long timeout, int displayId) {
+        public void waitForAllWindowsDrawn(Message message, long timeout, int displayId) {
+            Objects.requireNonNull(message.getTarget());
             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();
+                message.sendToTarget();
                 return;
             }
             boolean allWindowsDrawn = false;
@@ -7834,13 +7823,13 @@
                         }
                     }
 
-                    mWaitingForDrawnCallbacks.put(container, callback);
+                    mWaitingForDrawnCallbacks.put(container, message);
                     mH.sendNewMessageDelayed(H.WAITING_FOR_DRAWN_TIMEOUT, container, timeout);
                     checkDrawnWindowsLocked();
                 }
             }
             if (allWindowsDrawn) {
-                callback.run();
+                message.sendToTarget();
             }
         }