Hold Display Suspend blocker when going to Doze

The transition to Doze is asynchronous from the release of the display
suspend blocker.  This created a race condition where we could
potentially suspend before Doze service had a chance to transition to
Doze. This change holds the Display Suspend blocker until Doze service
acquires the DOZE lock.

Bug: 138828701
Test: atest PowerManagerService
Test: Manual, verify no repro of the bug.  Additional verify the device
does release the Display suspend blocker once we are in DOZE.

Change-Id: I777bc7963b1ab445378164b0f4d3f79113510b81
diff --git a/services/core/java/com/android/server/power/PowerManagerService.java b/services/core/java/com/android/server/power/PowerManagerService.java
index aa49ba6..e1b3e4d 100644
--- a/services/core/java/com/android/server/power/PowerManagerService.java
+++ b/services/core/java/com/android/server/power/PowerManagerService.java
@@ -542,6 +542,10 @@
     // True if we in the process of performing a forceSuspend
     private boolean mForceSuspendActive;
 
+    // Transition to Doze is in progress.  We have transitioned to WAKEFULNESS_DOZING,
+    // but the DreamService has not yet been told to start (it's an async process).
+    private boolean mDozeStartInProgress;
+
     private final class ForegroundProfileObserver extends SynchronousUserSwitchObserver {
         @Override
         public void onUserSwitching(int newUserId) throws RemoteException {}
@@ -1514,6 +1518,7 @@
             mLastSleepTime = eventTime;
             mLastSleepReason = reason;
             mSandmanSummoned = true;
+            mDozeStartInProgress = true;
             setWakefulnessLocked(WAKEFULNESS_DOZING, reason, eventTime);
 
             // Report the number of wake locks that will be cleared by going to sleep.
@@ -1601,6 +1606,10 @@
             mWakefulness = wakefulness;
             mWakefulnessChanging = true;
             mDirty |= DIRTY_WAKEFULNESS;
+
+            // This is only valid while we are in wakefulness dozing. Set to false otherwise.
+            mDozeStartInProgress &= (mWakefulness == WAKEFULNESS_DOZING);
+
             if (mNotifier != null) {
                 mNotifier.onWakefulnessChangeStarted(wakefulness, reason, eventTime);
             }
@@ -1631,6 +1640,9 @@
             if (mWakefulness == WAKEFULNESS_DOZING
                     && (mWakeLockSummary & WAKE_LOCK_DOZE) == 0) {
                 return; // wait until dream has enabled dozing
+            } else {
+                // Doze wakelock acquired (doze started) or device is no longer dozing.
+                mDozeStartInProgress = false;
             }
             if (mWakefulness == WAKEFULNESS_DOZING || mWakefulness == WAKEFULNESS_ASLEEP) {
                 logSleepTimeoutRecapturedLocked();
@@ -2309,6 +2321,10 @@
             isDreaming = false;
         }
 
+        // At this point, we either attempted to start the dream or no attempt will be made,
+        // so stop holding the display suspend blocker for Doze.
+        mDozeStartInProgress = false;
+
         // Update dream state.
         synchronized (mLock) {
             // Remember the initial battery level when the dream started.
@@ -2735,6 +2751,16 @@
         if (mScreenBrightnessBoostInProgress) {
             return true;
         }
+
+        // When we transition to DOZING, we have to keep the display suspend blocker
+        // up until the Doze service has a change to acquire the DOZE wakelock.
+        // Here we wait for mWakefulnessChanging to become false since the wakefulness
+        // transition to DOZING isn't considered "changed" until the doze wake lock is
+        // acquired.
+        if (mWakefulness == WAKEFULNESS_DOZING && mDozeStartInProgress) {
+            return true;
+        }
+
         // Let the system suspend if the screen is off or dozing.
         return false;
     }