Fix setCanTurnScreenOn for a launching activity

Originally the flag set by Activity#setCanTurnScreenOn takes effect
when the activity will be resumed. But if the activity uses the method
while the creation of activity is needed, the flag doesn't have the
chance to be set before the check for applying it because the client
side creation is after the launching flow.

By adding the check for ActivityRecord#canTurnScreenOn at the same
path of FLAG_TURN_SCREEN_ON, it is able to turn the screen on when the
activity relayouts to be visible.

AppWindowToken#canTurnScreenOn is renamed to currentLaunchCanTurnScreenOn
because it only applies to the current launch and avoids confusion with
the one in ActivityRecord.

Bug: 136214822
Test: atest WindowStateTests#testPrepareWindowToDisplayDuringRelayout
Test: atest ActivityVisibilityTests#testTurnScreenOnActivity
Change-Id: I9cb877357fb55c95b0d95044fe9e8ccba0c083ad
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 03cae42..48da52b 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -258,9 +258,9 @@
     ActivityRecord mActivityRecord;
 
     /**
-     * See {@link #canTurnScreenOn()}
+     * @see #currentLaunchCanTurnScreenOn()
      */
-    private boolean mCanTurnScreenOn = true;
+    private boolean mCurrentLaunchCanTurnScreenOn = true;
 
     /**
      * If we are running an animation, this determines the transition type. Must be one of
@@ -985,7 +985,7 @@
                 + " " + this);
         mAppStopped = false;
         // Allow the window to turn the screen on once the app is resumed again.
-        setCanTurnScreenOn(true);
+        setCurrentLaunchCanTurnScreenOn(true);
         if (!wasStopped) {
             destroySurfaces(true /*cleanupOnResume*/);
         }
@@ -2403,21 +2403,25 @@
     }
 
     /**
-     * Sets whether the current launch can turn the screen on. See {@link #canTurnScreenOn()}
+     * Sets whether the current launch can turn the screen on.
+     * @see #currentLaunchCanTurnScreenOn()
      */
-    void setCanTurnScreenOn(boolean canTurnScreenOn) {
-        mCanTurnScreenOn = canTurnScreenOn;
+    void setCurrentLaunchCanTurnScreenOn(boolean currentLaunchCanTurnScreenOn) {
+        mCurrentLaunchCanTurnScreenOn = currentLaunchCanTurnScreenOn;
     }
 
     /**
      * Indicates whether the current launch can turn the screen on. This is to prevent multiple
      * relayouts from turning the screen back on. The screen should only turn on at most
      * once per activity resume.
+     * <p>
+     * Note this flag is only meaningful when {@link WindowManager.LayoutParams#FLAG_TURN_SCREEN_ON}
+     * or {@link ActivityRecord#canTurnScreenOn} is set.
      *
-     * @return true if the screen can be turned on.
+     * @return {@code true} if the activity is ready to turn on the screen.
      */
-    boolean canTurnScreenOn() {
-        return mCanTurnScreenOn;
+    boolean currentLaunchCanTurnScreenOn() {
+        return mCurrentLaunchCanTurnScreenOn;
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index bf874be..6705f14 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2381,10 +2381,11 @@
 
     void prepareWindowToDisplayDuringRelayout(boolean wasVisible) {
         // We need to turn on screen regardless of visibility.
-        boolean hasTurnScreenOnFlag = (mAttrs.flags & FLAG_TURN_SCREEN_ON) != 0;
+        final boolean hasTurnScreenOnFlag = (mAttrs.flags & FLAG_TURN_SCREEN_ON) != 0
+                || (mAppToken != null && mAppToken.mActivityRecord.canTurnScreenOn());
 
         // The screen will turn on if the following conditions are met
-        // 1. The window has the flag FLAG_TURN_SCREEN_ON
+        // 1. The window has the flag FLAG_TURN_SCREEN_ON or ActivityRecord#canTurnScreenOn.
         // 2. The WMS allows theater mode.
         // 3. No AWT or the AWT allows the screen to be turned on. This should only be true once
         // per resume to prevent the screen getting getting turned on for each relayout. Set
@@ -2398,7 +2399,7 @@
             boolean allowTheaterMode = mWmService.mAllowTheaterModeWakeFromLayout
                     || Settings.Global.getInt(mWmService.mContext.getContentResolver(),
                             Settings.Global.THEATER_MODE_ON, 0) == 0;
-            boolean canTurnScreenOn = mAppToken == null || mAppToken.canTurnScreenOn();
+            boolean canTurnScreenOn = mAppToken == null || mAppToken.currentLaunchCanTurnScreenOn();
 
             if (allowTheaterMode && canTurnScreenOn && !mPowerManagerWrapper.isInteractive()) {
                 if (DEBUG_VISIBILITY || DEBUG_POWER) {
@@ -2409,7 +2410,7 @@
             }
 
             if (mAppToken != null) {
-                mAppToken.setCanTurnScreenOn(false);
+                mAppToken.setCurrentLaunchCanTurnScreenOn(false);
             }
         }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
index 1b57c79..36698ea 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
@@ -302,43 +302,38 @@
 
     @Test
     public void testPrepareWindowToDisplayDuringRelayout() {
-        testPrepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
-        testPrepareWindowToDisplayDuringRelayout(true /*wasVisible*/);
-
-        // Call prepareWindowToDisplayDuringRelayout for a window without FLAG_TURN_SCREEN_ON
-        // before calling prepareWindowToDisplayDuringRelayout for windows with flag in the same
-        // appWindowToken.
+        // Call prepareWindowToDisplayDuringRelayout for a window without FLAG_TURN_SCREEN_ON before
+        // calling setCurrentLaunchCanTurnScreenOn for windows with flag in the same appWindowToken.
         final AppWindowToken appWindowToken = createAppWindowToken(mDisplayContent,
                 WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_STANDARD);
         final WindowState first = createWindow(null, TYPE_APPLICATION, appWindowToken, "first");
         final WindowState second = createWindow(null, TYPE_APPLICATION, appWindowToken, "second");
         second.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
 
-        reset(sPowerManagerWrapper);
-        first.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
-        verify(sPowerManagerWrapper, never()).wakeUp(anyLong(), anyInt(), anyString());
-        assertTrue(appWindowToken.canTurnScreenOn());
-
-        reset(sPowerManagerWrapper);
-        second.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
-        verify(sPowerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
-        assertFalse(appWindowToken.canTurnScreenOn());
+        testPrepareWindowToDisplayDuringRelayout(first, false /* expectedWakeupCalled */,
+                true /* expectedCurrentLaunchCanTurnScreenOn */);
+        testPrepareWindowToDisplayDuringRelayout(second, true /* expectedWakeupCalled */,
+                false /* expectedCurrentLaunchCanTurnScreenOn */);
 
         // Call prepareWindowToDisplayDuringRelayout for two window that have FLAG_TURN_SCREEN_ON
         // from the same appWindowToken. Only one should trigger the wakeup.
-        appWindowToken.setCanTurnScreenOn(true);
+        appWindowToken.setCurrentLaunchCanTurnScreenOn(true);
         first.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
         second.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
 
-        reset(sPowerManagerWrapper);
-        first.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
-        verify(sPowerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
-        assertFalse(appWindowToken.canTurnScreenOn());
+        testPrepareWindowToDisplayDuringRelayout(first, true /* expectedWakeupCalled */,
+                false /* expectedCurrentLaunchCanTurnScreenOn */);
+        testPrepareWindowToDisplayDuringRelayout(second, false /* expectedWakeupCalled */,
+                false /* expectedCurrentLaunchCanTurnScreenOn */);
 
-        reset(sPowerManagerWrapper);
-        second.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
-        verify(sPowerManagerWrapper, never()).wakeUp(anyLong(), anyInt(), anyString());
-        assertFalse(appWindowToken.canTurnScreenOn());
+        // Without window flags, the state of ActivityRecord.canTurnScreenOn should still be able to
+        // turn on the screen.
+        appWindowToken.setCurrentLaunchCanTurnScreenOn(true);
+        first.mAttrs.flags &= ~WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
+        doReturn(true).when(appWindowToken.mActivityRecord).canTurnScreenOn();
+
+        testPrepareWindowToDisplayDuringRelayout(first, true /* expectedWakeupCalled */,
+                false /* expectedCurrentLaunchCanTurnScreenOn */);
 
         // Call prepareWindowToDisplayDuringRelayout for a windows that are not children of an
         // appWindowToken. Both windows have the FLAG_TURNS_SCREEN_ON so both should call wakeup
@@ -360,6 +355,22 @@
         verify(sPowerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
     }
 
+    private void testPrepareWindowToDisplayDuringRelayout(WindowState appWindow,
+            boolean expectedWakeupCalled, boolean expectedCurrentLaunchCanTurnScreenOn) {
+        reset(sPowerManagerWrapper);
+        appWindow.prepareWindowToDisplayDuringRelayout(false /* wasVisible */);
+
+        if (expectedWakeupCalled) {
+            verify(sPowerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
+        } else {
+            verify(sPowerManagerWrapper, never()).wakeUp(anyLong(), anyInt(), anyString());
+        }
+        // If wakeup is expected to be called, the currentLaunchCanTurnScreenOn should be false
+        // because the state will be consumed.
+        assertThat(appWindow.mAppToken.currentLaunchCanTurnScreenOn(),
+                is(expectedCurrentLaunchCanTurnScreenOn));
+    }
+
     @Test
     public void testCanAffectSystemUiFlags() {
         final WindowState app = createWindow(null, TYPE_APPLICATION, "app");
@@ -487,15 +498,6 @@
         assertThat(app.getWmDisplayCutout().getDisplayCutout(), is(cutout.inset(7, 10, 5, 20)));
     }
 
-    private void testPrepareWindowToDisplayDuringRelayout(boolean wasVisible) {
-        reset(sPowerManagerWrapper);
-        final WindowState root = createWindow(null, TYPE_APPLICATION, "root");
-        root.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
-
-        root.prepareWindowToDisplayDuringRelayout(wasVisible /*wasVisible*/);
-        verify(sPowerManagerWrapper).wakeUp(anyLong(), anyInt(), anyString());
-    }
-
     @Test
     public void testVisibilityChangeSwitchUser() {
         final WindowState window = createWindow(null, TYPE_APPLICATION, "app");