Merge "[waitForAllWindowsDrawn] Wait for transitions only on default displays" into main
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 87c0084..d8df645 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -7999,7 +7999,8 @@
             }
             boolean allWindowsDrawn = false;
             synchronized (mGlobalLock) {
-                if (mRoot.getDefaultDisplay().mDisplayUpdater.waitForTransition(message)) {
+                if (displayId == DEFAULT_DISPLAY
+                        && mRoot.getDefaultDisplay().mDisplayUpdater.waitForTransition(message)) {
                     // Use the ready-to-play of transition as the signal.
                     return;
                 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentDeferredUpdateTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentDeferredUpdateTests.java
index f843386..affaad6 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentDeferredUpdateTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentDeferredUpdateTests.java
@@ -16,12 +16,16 @@
 
 package com.android.server.wm;
 
+import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
+
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.any;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.never;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static com.android.server.wm.WindowStateAnimator.HAS_DRAWN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -37,6 +41,7 @@
 
 import androidx.test.filters.SmallTest;
 
+import com.android.server.LocalServices;
 import com.android.server.wm.TransitionController.OnStartCollect;
 import com.android.window.flags.Flags;
 
@@ -57,18 +62,29 @@
 public class DisplayContentDeferredUpdateTests extends WindowTestsBase {
 
     // The fields to override the current DisplayInfo.
-    private String mUniqueId;
+    private String mUniqueId = "initial_unique_id";
+    private String mSecondaryUniqueId = "secondary_initial_unique_id";
     private int mColorMode;
     private int mLogicalDensityDpi;
 
+    private DisplayContent mSecondaryDisplayContent;
+
     private final Message mScreenUnblocker = mock(Message.class);
+    private final Message mSecondaryScreenUnblocker = mock(Message.class);
+
+    private WindowManagerInternal mWmInternal;
 
     @Before
     public void before() {
+        when(mScreenUnblocker.getTarget()).thenReturn(mWm.mH);
         doReturn(true).when(mDisplayContent).getLastHasContent();
-        mockTransitionsController(/* enabled= */ true);
-        mockRemoteDisplayChangeController();
-        performInitialDisplayUpdate();
+
+        mockTransitionsController();
+
+        mockRemoteDisplayChangeController(mDisplayContent);
+        performInitialDisplayUpdate(mDisplayContent);
+
+        mWmInternal = LocalServices.getService(WindowManagerInternal.class);
     }
 
     @Test
@@ -245,24 +261,90 @@
         verify(mScreenUnblocker, never()).sendToTarget();
     }
 
-    private void mockTransitionsController(boolean enabled) {
-        spyOn(mDisplayContent.mTransitionController);
-        when(mDisplayContent.mTransitionController.isShellTransitionsEnabled()).thenReturn(enabled);
-        doReturn(true).when(mDisplayContent.mTransitionController).startCollectOrQueue(any(),
-                any());
+    @Test
+    public void testTwoDisplayUpdateAtTheSameTime_bothDisplaysAreUnblocked() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_WAIT_FOR_TRANSITION_ON_DISPLAY_SWITCH);
+        prepareSecondaryDisplay();
+
+        final WindowState defaultDisplayWindow = createWindow(/* parent= */ null,
+                TYPE_BASE_APPLICATION, mDisplayContent, "DefaultDisplayWindow");
+        final WindowState secondaryDisplayWindow = createWindow(/* parent= */ null,
+                TYPE_BASE_APPLICATION, mSecondaryDisplayContent, "SecondaryDisplayWindow");
+        makeWindowVisibleAndNotDrawn(defaultDisplayWindow, secondaryDisplayWindow);
+
+        // Mark as display switching only for the default display as we filter out
+        // non-default display switching events in the display policy
+        mDisplayContent.mDisplayUpdater.onDisplaySwitching(/* switching= */ true);
+
+        mWmInternal.waitForAllWindowsDrawn(mScreenUnblocker,
+                /* timeout= */ Integer.MAX_VALUE, DEFAULT_DISPLAY);
+        mWmInternal.waitForAllWindowsDrawn(mSecondaryScreenUnblocker,
+                /* timeout= */ Integer.MAX_VALUE, mSecondaryDisplayContent.getDisplayId());
+
+        // Perform display update for both displays at the same time
+        mUniqueId = "new_default_display_unique_id";
+        mSecondaryUniqueId = "new_secondary_display_unique_id";
+        mDisplayContent.requestDisplayUpdate(mock(Runnable.class));
+        mSecondaryDisplayContent.requestDisplayUpdate(mock(Runnable.class));
+
+        when(mDisplayContent.mTransitionController.inTransition()).thenReturn(true);
+
+        // Notify that both transitions started collecting
+        captureStartTransitionCollection().getAllValues().forEach((callback) ->
+                callback.onCollectStarted(/* deferred= */ true));
+
+        // Verify that screens are not unblocked yet
+        verify(mScreenUnblocker, never()).sendToTarget();
+        verify(mSecondaryScreenUnblocker, never()).sendToTarget();
+
+        // Make all secondary display windows drawn
+        secondaryDisplayWindow.mWinAnimator.mDrawState = HAS_DRAWN;
+        mWm.mRoot.performSurfacePlacement();
+
+        // Verify that only secondary screen is unblocked as it uses
+        // the legacy waitForAllWindowsDrawn path
+        verify(mScreenUnblocker, never()).sendToTarget();
+        verify(mSecondaryScreenUnblocker).sendToTarget();
+
+        // Mark start transactions as presented
+        when(mDisplayContent.mTransitionController.inTransition()).thenReturn(false);
+        captureRequestedTransition().getAllValues().forEach(
+                this::makeTransitionTransactionCompleted);
+
+        // Verify that the default screen unblocker is sent only after start transaction
+        // of the Shell transition is presented
+        verify(mScreenUnblocker).sendToTarget();
     }
 
-    private void mockRemoteDisplayChangeController() {
-        spyOn(mDisplayContent.mRemoteDisplayChangeController);
-        doReturn(true).when(mDisplayContent.mRemoteDisplayChangeController)
+    private void prepareSecondaryDisplay() {
+        mSecondaryDisplayContent = createNewDisplay();
+        when(mSecondaryScreenUnblocker.getTarget()).thenReturn(mWm.mH);
+        doReturn(true).when(mSecondaryDisplayContent).getLastHasContent();
+        mockRemoteDisplayChangeController(mSecondaryDisplayContent);
+        performInitialDisplayUpdate(mSecondaryDisplayContent);
+    }
+
+    private void mockTransitionsController() {
+        spyOn(mDisplayContent.mTransitionController);
+        when(mDisplayContent.mTransitionController.isShellTransitionsEnabled())
+                .thenReturn(true);
+        doReturn(mock(Transition.class)).when(mDisplayContent.mTransitionController)
+                .createTransition(anyInt(), anyInt());
+        doReturn(true).when(mDisplayContent.mTransitionController)
+                .startCollectOrQueue(any(), any());
+    }
+
+    private void mockRemoteDisplayChangeController(DisplayContent displayContent) {
+        spyOn(displayContent.mRemoteDisplayChangeController);
+        doReturn(true).when(displayContent.mRemoteDisplayChangeController)
                 .performRemoteDisplayChange(anyInt(), anyInt(), any(), any());
     }
 
     private ArgumentCaptor<OnStartCollect> captureStartTransitionCollection() {
         ArgumentCaptor<OnStartCollect> callbackCaptor =
                 ArgumentCaptor.forClass(OnStartCollect.class);
-        verify(mDisplayContent.mTransitionController, atLeast(1)).startCollectOrQueue(any(),
-                callbackCaptor.capture());
+        verify(mDisplayContent.mTransitionController, atLeast(1))
+                .startCollectOrQueue(any(), callbackCaptor.capture());
         return callbackCaptor;
     }
 
@@ -283,20 +365,23 @@
         }
     }
 
-    private void performInitialDisplayUpdate() {
-        mUniqueId = "initial_unique_id";
+    private void performInitialDisplayUpdate(DisplayContent displayContent) {
         mColorMode = 0;
         mLogicalDensityDpi = 400;
 
-        spyOn(mDisplayContent.mDisplay);
+        spyOn(displayContent.mDisplay);
         doAnswer(invocation -> {
             DisplayInfo info = invocation.getArgument(0);
-            info.uniqueId = mUniqueId;
+            if (displayContent.isDefaultDisplay) {
+                info.uniqueId = mUniqueId;
+            } else {
+                info.uniqueId = mSecondaryUniqueId;
+            }
             info.colorMode = mColorMode;
             info.logicalDensityDpi = mLogicalDensityDpi;
             return null;
-        }).when(mDisplayContent.mDisplay).getDisplayInfo(any());
+        }).when(displayContent.mDisplay).getDisplayInfo(any());
         Runnable onUpdated = mock(Runnable.class);
-        mDisplayContent.requestDisplayUpdate(onUpdated);
+        displayContent.requestDisplayUpdate(onUpdated);
     }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index 4a8a2e6..bcf4ebc 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -55,6 +55,7 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.server.wm.WindowContainer.POSITION_BOTTOM;
 import static com.android.server.wm.WindowContainer.POSITION_TOP;
+import static com.android.server.wm.WindowStateAnimator.DRAW_PENDING;
 import static com.android.server.wm.WindowStateAnimator.HAS_DRAWN;
 
 import static org.junit.Assert.assertEquals;
@@ -693,6 +694,13 @@
         }
     }
 
+    static void makeWindowVisibleAndNotDrawn(WindowState... windows) {
+        makeWindowVisible(windows);
+        for (WindowState win : windows) {
+            win.mWinAnimator.mDrawState = DRAW_PENDING;
+        }
+    }
+
     static void makeLastConfigReportedToClient(WindowState w, boolean visible) {
         w.fillClientWindowFramesAndConfiguration(new ClientWindowFrames(),
                 new MergedConfiguration(), new ActivityWindowInfo(), true /* useLatestConfig */,