DisplayManagerInternal: add cancelBlockScreenOn for DisplayOffloadSession

- Add a cancelBlockScreenOn method in DisplayOffloadSession
- Add test
- Minor test fixes

Flag: com.android.server.display.feature.flags.offload_session_cancel_block_screen_on
Test: atest DisplayOffloadSessionImplTest
Test: atest DisplayPowerControllerTest
Bug: 346420374
Bug: 347415946
Change-Id: Ic043208968bf516a58420e1e293c663d44061782
diff --git a/core/java/android/hardware/display/DisplayManagerInternal.java b/core/java/android/hardware/display/DisplayManagerInternal.java
index b2dcf90..91caedc 100644
--- a/core/java/android/hardware/display/DisplayManagerInternal.java
+++ b/core/java/android/hardware/display/DisplayManagerInternal.java
@@ -740,6 +740,12 @@
          */
         void onBlockingScreenOn(Runnable unblocker);
 
+        /**
+         * Called while display is turning to screen state other than state ON to notify that any
+         * pending work from the previous blockScreenOn call should have been cancelled.
+         */
+        void cancelBlockScreenOn();
+
         /** Whether auto brightness update in doze is allowed */
         boolean allowAutoBrightnessInDoze();
     }
@@ -774,6 +780,12 @@
         boolean blockScreenOn(Runnable unblocker);
 
         /**
+         * Called while display is turning to screen state other than state ON to notify that any
+         * pending work from the previous blockScreenOn call should have been cancelled.
+         */
+        void cancelBlockScreenOn();
+
+        /**
          * Get the brightness levels used to determine automatic brightness based on lux levels.
          * @param mode The auto-brightness mode
          *             (AutomaticBrightnessController.AutomaticBrightnessMode)
diff --git a/services/core/java/com/android/server/display/DisplayOffloadSessionImpl.java b/services/core/java/com/android/server/display/DisplayOffloadSessionImpl.java
index 0fef55d..a188e79 100644
--- a/services/core/java/com/android/server/display/DisplayOffloadSessionImpl.java
+++ b/services/core/java/com/android/server/display/DisplayOffloadSessionImpl.java
@@ -84,6 +84,14 @@
     }
 
     @Override
+    public void cancelBlockScreenOn() {
+        if (mDisplayOffloader == null) {
+            return;
+        }
+        mDisplayOffloader.cancelBlockScreenOn();
+    }
+
+    @Override
     public float[] getAutoBrightnessLevels(int mode) {
         if (mode < 0 || mode > AUTO_BRIGHTNESS_MODE_MAX) {
             throw new IllegalArgumentException("Unknown auto-brightness mode: " + mode);
diff --git a/services/core/java/com/android/server/display/DisplayPowerController.java b/services/core/java/com/android/server/display/DisplayPowerController.java
index 65a729a..f857167 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController.java
@@ -2108,6 +2108,17 @@
                 Trace.TRACE_TAG_POWER, SCREEN_ON_BLOCKED_BY_DISPLAYOFFLOAD_TRACE_NAME, 0);
     }
 
+    private void cancelUnblockScreenOnByDisplayOffload() {
+        if (mDisplayOffloadSession == null) {
+            return;
+        }
+        if (mPendingScreenOnUnblockerByDisplayOffload != null) {
+            // Already unblocked.
+            return;
+        }
+        mDisplayOffloadSession.cancelBlockScreenOn();
+    }
+
     private boolean setScreenState(int state, @Display.StateReason int reason) {
         return setScreenState(state, reason, false /*reportOnly*/);
     }
@@ -2124,6 +2135,9 @@
             blockScreenOnByDisplayOffload(mDisplayOffloadSession);
         } else if (!isOn && mScreenTurningOnWasBlockedByDisplayOffload) {
             // No longer turning screen on, so unblock previous screen on blocking immediately.
+            if (mFlags.isOffloadSessionCancelBlockScreenOnEnabled()) {
+                cancelUnblockScreenOnByDisplayOffload();
+            }
             unblockScreenOnByDisplayOffload();
             mScreenTurningOnWasBlockedByDisplayOffload = false;
         }
diff --git a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
index f56d803..41d18cd 100644
--- a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
+++ b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
@@ -179,6 +179,11 @@
             Flags::offloadDozeOverrideHoldsWakelock
     );
 
+    private final FlagState mOffloadSessionCancelBlockScreenOn =
+            new FlagState(
+                    Flags.FLAG_OFFLOAD_SESSION_CANCEL_BLOCK_SCREEN_ON,
+                    Flags::offloadSessionCancelBlockScreenOn);
+
     /**
      * @return {@code true} if 'port' is allowed in display layout configuration file.
      */
@@ -352,6 +357,10 @@
         return mOffloadDozeOverrideHoldsWakelock.isEnabled();
     }
 
+    public boolean isOffloadSessionCancelBlockScreenOnEnabled() {
+        return mOffloadSessionCancelBlockScreenOn.isEnabled();
+    }
+
     /**
      * @return Whether to ignore preferredRefreshRate app request conversion to display mode or not
      */
@@ -399,6 +408,7 @@
         pw.println(" " + mIgnoreAppPreferredRefreshRate);
         pw.println(" " + mSynthetic60hzModes);
         pw.println(" " + mOffloadDozeOverrideHoldsWakelock);
+        pw.println(" " + mOffloadSessionCancelBlockScreenOn);
     }
 
     private static class FlagState {
diff --git a/services/core/java/com/android/server/display/feature/display_flags.aconfig b/services/core/java/com/android/server/display/feature/display_flags.aconfig
index 95d0ca3..1ea5c0b 100644
--- a/services/core/java/com/android/server/display/feature/display_flags.aconfig
+++ b/services/core/java/com/android/server/display/feature/display_flags.aconfig
@@ -299,3 +299,11 @@
       purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "offload_session_cancel_block_screen_on"
+    namespace: "wear_frameworks"
+    description: "Flag for DisplayPowerController to start notifying DisplayOffloadSession about cancelling screen on blocker."
+    bug: "331725519"
+    is_fixed_read_only: true
+}
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayOffloadSessionImplTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayOffloadSessionImplTest.java
index 4409051..30c384a 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayOffloadSessionImplTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayOffloadSessionImplTest.java
@@ -94,4 +94,11 @@
 
         verify(mDisplayOffloader).onBlockingScreenOn(eq(unblocker));
     }
+
+    @Test
+    public void testUnblockScreenOn() {
+        mSession.cancelBlockScreenOn();
+
+        verify(mDisplayOffloader).cancelBlockScreenOn();
+    }
 }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
index 95f0b65..bb774ee 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
@@ -105,7 +105,6 @@
 
 import java.util.List;
 
-
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public final class DisplayPowerControllerTest {
@@ -1660,6 +1659,8 @@
         int initState = Display.STATE_OFF;
         mHolder = createDisplayPowerController(DISPLAY_ID, UNIQUE_ID);
         mHolder.dpc.setDisplayOffloadSession(mDisplayOffloadSession);
+        when(mDisplayOffloadSession.blockScreenOn(any())).thenReturn(true);
+
         // start with OFF.
         when(mHolder.displayPowerState.getScreenState()).thenReturn(initState);
         DisplayPowerRequest dpr = new DisplayPowerRequest();
@@ -1673,6 +1674,7 @@
         advanceTime(1); // Run updatePowerState
 
         verify(mDisplayOffloadSession).blockScreenOn(any(Runnable.class));
+        verify(mDisplayOffloadSession, never()).cancelBlockScreenOn();
     }
 
     @Test
@@ -1680,6 +1682,8 @@
         // set up.
         int initState = Display.STATE_ON;
         mHolder.dpc.setDisplayOffloadSession(mDisplayOffloadSession);
+        when(mDisplayOffloadSession.blockScreenOn(any())).thenReturn(true);
+
         // start with ON.
         when(mHolder.displayPowerState.getScreenState()).thenReturn(initState);
         DisplayPowerRequest dpr = new DisplayPowerRequest();
@@ -1692,7 +1696,78 @@
         mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
         advanceTime(1); // Run updatePowerState
 
-        verify(mDisplayOffloadSession, never()).blockScreenOn(any(Runnable.class));
+        // No cancelBlockScreenOn call because we didn't block.
+        verify(mDisplayOffloadSession, never()).cancelBlockScreenOn();
+    }
+
+    @RequiresFlagsEnabled(Flags.FLAG_OFFLOAD_SESSION_CANCEL_BLOCK_SCREEN_ON)
+    @Test
+    public void testOffloadBlocker_turnON_thenOFF_cancelBlockScreenOnNotCalledIfUnblocked() {
+        // Set up.
+        int initState = Display.STATE_OFF;
+        mHolder = createDisplayPowerController(DISPLAY_ID, UNIQUE_ID);
+        mHolder.dpc.setDisplayOffloadSession(mDisplayOffloadSession);
+        when(mDisplayOffloadSession.blockScreenOn(any())).thenReturn(true);
+
+        // Start with OFF.
+        when(mHolder.displayPowerState.getScreenState()).thenReturn(initState);
+        DisplayPowerRequest dpr = new DisplayPowerRequest();
+        dpr.policy = DisplayPowerRequest.POLICY_OFF;
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        // Go to ON.
+        dpr.policy = DisplayPowerRequest.POLICY_BRIGHT;
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        ArgumentCaptor<Runnable> argumentCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mDisplayOffloadSession).blockScreenOn(argumentCaptor.capture());
+
+        // Unblocked
+        argumentCaptor.getValue().run();
+        advanceTime(1); // Run updatePowerState
+
+        // Go to OFF immediately
+        dpr.policy = DisplayPowerRequest.POLICY_OFF;
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        // No cancelBlockScreenOn call because we already unblocked
+        verify(mDisplayOffloadSession, never()).cancelBlockScreenOn();
+    }
+
+    @RequiresFlagsEnabled(Flags.FLAG_OFFLOAD_SESSION_CANCEL_BLOCK_SCREEN_ON)
+    @Test
+    public void testOffloadBlocker_turnON_thenOFF_cancelBlockScreenOn() {
+        // Set up.
+        int initState = Display.STATE_OFF;
+        mHolder = createDisplayPowerController(DISPLAY_ID, UNIQUE_ID);
+        mHolder.dpc.setDisplayOffloadSession(mDisplayOffloadSession);
+        when(mDisplayOffloadSession.blockScreenOn(any())).thenReturn(true);
+
+        // Start with OFF.
+        when(mHolder.displayPowerState.getScreenState()).thenReturn(initState);
+        DisplayPowerRequest dpr = new DisplayPowerRequest();
+        dpr.policy = DisplayPowerRequest.POLICY_OFF;
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        // Go to ON.
+        dpr.policy = DisplayPowerRequest.POLICY_BRIGHT;
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        // We should call blockScreenOn
+        verify(mDisplayOffloadSession).blockScreenOn(any(Runnable.class));
+
+        // Go to OFF immediately
+        dpr.policy = DisplayPowerRequest.POLICY_OFF;
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        // We should call cancelBlockScreenOn
+        verify(mDisplayOffloadSession).cancelBlockScreenOn();
     }
 
     @Test
diff --git a/services/tests/displayservicetests/src/com/android/server/display/LocalDisplayAdapterTest.java b/services/tests/displayservicetests/src/com/android/server/display/LocalDisplayAdapterTest.java
index 01ff35f..a7e0ebd 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/LocalDisplayAdapterTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/LocalDisplayAdapterTest.java
@@ -33,7 +33,6 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.when;
 
@@ -83,7 +82,6 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
-
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class LocalDisplayAdapterTest {
@@ -126,7 +124,7 @@
 
     private DisplayOffloadSessionImpl mDisplayOffloadSession;
 
-    private DisplayOffloader mDisplayOffloader;
+    @Mock DisplayOffloader mDisplayOffloader;
 
     private TestListener mListener = new TestListener();
 
@@ -1249,24 +1247,8 @@
     }
 
     private void initDisplayOffloadSession() {
-        mDisplayOffloader = spy(new DisplayOffloader() {
-            @Override
-            public boolean startOffload() {
-                return true;
-            }
-
-            @Override
-            public void stopOffload() {}
-
-            @Override
-            public void onBlockingScreenOn(Runnable unblocker) {}
-
-            @Override
-            public boolean allowAutoBrightnessInDoze() {
-                return true;
-            }
-        });
-
+        when(mDisplayOffloader.startOffload()).thenReturn(true);
+        when(mDisplayOffloader.allowAutoBrightnessInDoze()).thenReturn(true);
         mDisplayOffloadSession = new DisplayOffloadSessionImpl(mDisplayOffloader,
                 mMockedDisplayPowerController);
     }