diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java b/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
index c633aa15c8e34..951b27f75aeb6 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
@@ -225,8 +225,10 @@ public class DozeMachine {
         boolean newPolicy = wakeLockPolicy(newState);
         if (mWakeLockHeldForCurrentState && !newPolicy) {
             mWakeLock.release();
+            mWakeLockHeldForCurrentState = false;
         } else if (!mWakeLockHeldForCurrentState && newPolicy) {
             mWakeLock.acquire();
+            mWakeLockHeldForCurrentState = true;
         }
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
index ba7c923efe711..863f0e5a9897a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
@@ -187,6 +187,19 @@ public class DozeMachineTest {
         assertFalse(mWakeLockFake.isHeld());
     }
 
+    @Test
+    @UiThreadTest
+    public void testWakeLock_releasedAfterPulse() {
+        mMachine.requestState(INITIALIZED);
+
+        mMachine.requestState(DOZE);
+        mMachine.requestState(DOZE_REQUEST_PULSE);
+        mMachine.requestState(DOZE_PULSING);
+        mMachine.requestState(DOZE_PULSE_DONE);
+
+        assertFalse(mWakeLockFake.isHeld());
+    }
+
     @Test
     @UiThreadTest
     public void testScreen_offInDoze() {