From 39880be88c3b40c45e189199d7b7a633410bd85d Mon Sep 17 00:00:00 2001
From: Adrian Roos <roosa@google.com>
Date: Mon, 23 Apr 2018 17:41:17 +0200
Subject: [PATCH] AOD: Fix WakeLock leak

Bug: 78402666
Test: Toggle screen off and on again real quick. Verify "Doze" wake lock was not leaked.
Change-Id: Ie34f8aeeb5d10628f2a40567267ab36e694e21bb
---
 .../systemui/doze/DozeScreenState.java        | 17 +++++------
 .../systemui/doze/DozeScreenStateTest.java    | 30 +++++++++++++++----
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
index f72bff5d651ed..a0fdb9dc5b90f 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
@@ -21,6 +21,7 @@ import android.util.Log;
 import android.view.Display;
 
 import com.android.systemui.statusbar.phone.DozeParameters;
+import com.android.systemui.util.wakelock.SettableWakeLock;
 import com.android.systemui.util.wakelock.WakeLock;
 
 /**
@@ -43,15 +44,14 @@ public class DozeScreenState implements DozeMachine.Part {
     private final DozeParameters mParameters;
 
     private int mPendingScreenState = Display.STATE_UNKNOWN;
-    private boolean mWakeLockHeld;
-    private WakeLock mWakeLock;
+    private SettableWakeLock mWakeLock;
 
     public DozeScreenState(DozeMachine.Service service, Handler handler,
             DozeParameters parameters, WakeLock wakeLock) {
         mDozeService = service;
         mHandler = handler;
         mParameters = parameters;
-        mWakeLock = wakeLock;
+        mWakeLock = new SettableWakeLock(wakeLock);
     }
 
     @Override
@@ -64,6 +64,7 @@ public class DozeScreenState implements DozeMachine.Part {
             mHandler.removeCallbacks(mApplyPendingScreenState);
 
             applyScreenState(screenState);
+            mWakeLock.setAcquired(false);
             return;
         }
 
@@ -84,9 +85,8 @@ public class DozeScreenState implements DozeMachine.Part {
             boolean shouldDelayTransition = newState == DozeMachine.State.DOZE_AOD
                     && mParameters.shouldControlScreenOff();
 
-            if (!mWakeLockHeld && shouldDelayTransition) {
-                mWakeLockHeld = true;
-                mWakeLock.acquire();
+            if (shouldDelayTransition) {
+                mWakeLock.setAcquired(true);
             }
 
             if (!messagePending) {
@@ -118,10 +118,7 @@ public class DozeScreenState implements DozeMachine.Part {
             if (DEBUG) Log.d(TAG, "setDozeScreenState(" + screenState + ")");
             mDozeService.setDozeScreenState(screenState);
             mPendingScreenState = Display.STATE_UNKNOWN;
-            if (mWakeLockHeld) {
-                mWakeLockHeld = false;
-                mWakeLock.release();
-            }
+            mWakeLock.setAcquired(false);
         }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
index 66836365705e4..b6039b66e6cbf 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
@@ -25,8 +25,10 @@ import static com.android.systemui.doze.DozeMachine.State.INITIALIZED;
 import static com.android.systemui.doze.DozeMachine.State.UNINITIALIZED;
 import static com.android.systemui.utils.os.FakeHandler.Mode.QUEUEING;
 
+import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
@@ -41,6 +43,7 @@ import android.view.Display;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.phone.DozeParameters;
 import com.android.systemui.util.wakelock.WakeLock;
+import com.android.systemui.util.wakelock.WakeLockFake;
 import com.android.systemui.utils.os.FakeHandler;
 
 import org.junit.Before;
@@ -58,8 +61,7 @@ public class DozeScreenStateTest extends SysuiTestCase {
     FakeHandler mHandlerFake;
     @Mock
     DozeParameters mDozeParameters;
-    @Mock
-    WakeLock mWakeLock;
+    WakeLockFake mWakeLock;
 
     @Before
     public void setUp() throws Exception {
@@ -68,6 +70,7 @@ public class DozeScreenStateTest extends SysuiTestCase {
         when(mDozeParameters.getAlwaysOn()).thenReturn(true);
         mServiceFake = new DozeServiceFake();
         mHandlerFake = new FakeHandler(Looper.getMainLooper());
+        mWakeLock = new WakeLockFake();
         mScreen = new DozeScreenState(mServiceFake, mHandlerFake, mDozeParameters, mWakeLock);
     }
 
@@ -158,14 +161,29 @@ public class DozeScreenStateTest extends SysuiTestCase {
 
         mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
         mHandlerFake.dispatchQueuedMessages();
-        reset(mWakeLock);
 
         mScreen.transitionTo(INITIALIZED, DOZE_AOD);
-        verify(mWakeLock).acquire();
-        verify(mWakeLock, never()).release();
+        assertThat(mWakeLock.isHeld(), is(true));
 
         mHandlerFake.dispatchQueuedMessages();
-        verify(mWakeLock).release();
+        assertThat(mWakeLock.isHeld(), is(false));
+    }
+
+    @Test
+    public void test_releasesWakeLock_abortingLowPowerDelayed() {
+        // Transition to low power mode will be delayed to let
+        // animations play at 60 fps.
+        when(mDozeParameters.shouldControlScreenOff()).thenReturn(true);
+        mHandlerFake.setMode(QUEUEING);
+
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mHandlerFake.dispatchQueuedMessages();
+
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+        assertThat(mWakeLock.isHeld(), is(true));
+        mScreen.transitionTo(DOZE_AOD, FINISH);
+
+        assertThat(mWakeLock.isHeld(), is(false));
     }
 
 }
\ No newline at end of file