Move Wakelock acquisition and release logic to WakeLockController

This is to decouple the dependency of DisplayPowerController from
suspend blockers lifecycle. We do not expect any behavioural changes in
the functioning of a device post this.

Bug: 245920251
Test: atest WakelockControllerTest
Test: Manual with DisplayPowerController2
Change-Id: I115a31a0557e44d1cbe70362c6dd2b315ef1241d

4 files changed, 607 insertions, 151 deletions

diff --git a/services/core/java/com/android/server/display/DisplayPowerController2.java b/services/core/java/com/android/server/display/DisplayPowerController2.java
index c734095a4743..b8a881ed33d4 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController2.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController2.java
@@ -171,10 +171,6 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
     // Our handler.
     private final DisplayControllerHandler mHandler;
 
-    // Asynchronous callbacks into the power manager service.
-    // Only invoked from the handler thread while no locks are held.
-    private final DisplayPowerCallbacks mCallbacks;
-
     // Battery stats.
     @Nullable
     private final IBatteryStats mBatteryStats;
@@ -323,7 +319,10 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
 
     // The raw non-debounced proximity sensor state.
     private int mPendingProximity = PROXIMITY_UNKNOWN;
-    private long mPendingProximityDebounceTime = -1; // -1 if fully debounced
+
+    // -1 if fully debounced. Else, represents the time in ms when the debounce suspend blocker will
+    // be removed. Applies for both positive and negative proximity flips.
+    private long mPendingProximityDebounceTime = -1;
 
     // True if the screen was turned off because of the proximity sensor.
     // When the screen turns on again, we report user activity to the power manager.
@@ -339,9 +338,6 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
     // turning off the screen.
     private boolean mPendingScreenOff;
 
-    // True if we have unfinished business and are holding a suspend blocker.
-    private boolean mUnfinishedBusiness;
-
     // The elapsed real time when the screen on was blocked.
     private long mScreenOnBlockStartRealTime;
     private long mScreenOffBlockStartRealTime;
@@ -407,6 +403,10 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
     // Keeps a record of brightness changes for dumpsys.
     private RingBuffer<BrightnessEvent> mBrightnessEventRingBuffer;
 
+    // Controls and tracks all the wakelocks that are acquired/released by the system. Also acts as
+    // a medium of communication between this class and the PowerManagerService.
+    private final WakelockController mWakelockController;
+
     // A record of state for skipping brightness ramps.
     private int mSkipRampState = RAMP_STATE_SKIP_NONE;
 
@@ -467,18 +467,6 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
 
     private boolean mIsRbcActive;
 
-    // Whether there's a callback to tell listeners the display has changed scheduled to run. When
-    // true it implies a wakelock is being held to guarantee the update happens before we collapse
-    // into suspend and so needs to be cleaned up if the thread is exiting.
-    // Should only be accessed on the Handler thread.
-    private boolean mOnStateChangedPending;
-
-    // Count of proximity messages currently on this DPC's Handler. Used to keep track of how many
-    // suspend blocker acquisitions are pending when shutting down this DPC.
-    // Should only be accessed on the Handler thread.
-    private int mOnProximityPositiveMessages;
-    private int mOnProximityNegativeMessages;
-
     // Animators.
     private ObjectAnimator mColorFadeOnAnimator;
     private ObjectAnimator mColorFadeOffAnimator;
@@ -490,13 +478,6 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
 
     private DisplayDeviceConfig mDisplayDeviceConfig;
 
-    // Identifiers for suspend blocker acuisition requests
-    private final String mSuspendBlockerIdUnfinishedBusiness;
-    private final String mSuspendBlockerIdOnStateChanged;
-    private final String mSuspendBlockerIdProxPositive;
-    private final String mSuspendBlockerIdProxNegative;
-    private final String mSuspendBlockerIdProxDebounce;
-
     /**
      * Creates the display power controller.
      */
@@ -510,12 +491,8 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         mClock = mInjector.getClock();
         mLogicalDisplay = logicalDisplay;
         mDisplayId = mLogicalDisplay.getDisplayIdLocked();
+        mWakelockController = mInjector.getWakelockController(mDisplayId, callbacks);
         mTag = "DisplayPowerController2[" + mDisplayId + "]";
-        mSuspendBlockerIdUnfinishedBusiness = getSuspendBlockerUnfinishedBusinessId(mDisplayId);
-        mSuspendBlockerIdOnStateChanged = getSuspendBlockerOnStateChangedId(mDisplayId);
-        mSuspendBlockerIdProxPositive = getSuspendBlockerProxPositiveId(mDisplayId);
-        mSuspendBlockerIdProxNegative = getSuspendBlockerProxNegativeId(mDisplayId);
-        mSuspendBlockerIdProxDebounce = getSuspendBlockerProxDebounceId(mDisplayId);
 
         mDisplayDevice = mLogicalDisplay.getPrimaryDisplayDeviceLocked();
         mUniqueDisplayId = logicalDisplay.getPrimaryDisplayDeviceLocked().getUniqueId();
@@ -531,7 +508,6 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         }
 
         mSettingsObserver = new SettingsObserver(mHandler);
-        mCallbacks = callbacks;
         mSensorManager = sensorManager;
         mWindowManagerPolicy = LocalServices.getService(WindowManagerPolicy.class);
         mBlanker = blanker;
@@ -1143,22 +1119,10 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         mHandler.removeCallbacksAndMessages(null);
 
         // Release any outstanding wakelocks we're still holding because of pending messages.
-        if (mUnfinishedBusiness) {
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdUnfinishedBusiness);
-            mUnfinishedBusiness = false;
-        }
-        if (mOnStateChangedPending) {
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdOnStateChanged);
-            mOnStateChangedPending = false;
-        }
-        for (int i = 0; i < mOnProximityPositiveMessages; i++) {
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxPositive);
-        }
-        mOnProximityPositiveMessages = 0;
-        for (int i = 0; i < mOnProximityNegativeMessages; i++) {
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxNegative);
-        }
-        mOnProximityNegativeMessages = 0;
+        mWakelockController.releaseUnfinishedBusinessSuspendBlocker();
+        mWakelockController.releaseStateChangedSuspendBlocker();
+        mWakelockController.releaseProxPositiveSuspendBlocker();
+        mWakelockController.releaseProxNegativeSuspendBlocker();
 
         final float brightness = mPowerState != null
                 ? mPowerState.getScreenBrightness()
@@ -1748,12 +1712,8 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         }
 
         // Grab a wake lock if we have unfinished business.
-        if (!finished && !mUnfinishedBusiness) {
-            if (DEBUG) {
-                Slog.d(mTag, "Unfinished business...");
-            }
-            mCallbacks.acquireSuspendBlocker(mSuspendBlockerIdUnfinishedBusiness);
-            mUnfinishedBusiness = true;
+        if (!finished) {
+            mWakelockController.acquireUnfinishedBusinessSuspendBlocker();
         }
 
         // Notify the power manager when ready.
@@ -1772,12 +1732,8 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         }
 
         // Release the wake lock when we have no unfinished business.
-        if (finished && mUnfinishedBusiness) {
-            if (DEBUG) {
-                Slog.d(mTag, "Finished business...");
-            }
-            mUnfinishedBusiness = false;
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdUnfinishedBusiness);
+        if (finished) {
+            mWakelockController.releaseUnfinishedBusinessSuspendBlocker();
         }
 
         // Record if dozing for future comparison.
@@ -2274,7 +2230,12 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
                 mPendingProximity = PROXIMITY_UNKNOWN;
                 mHandler.removeMessages(MSG_PROXIMITY_SENSOR_DEBOUNCED);
                 mSensorManager.unregisterListener(mProximitySensorListener);
-                clearPendingProximityDebounceTime(); // release wake lock (must be last)
+                // release wake lock(must be last)
+                boolean proxDebounceSuspendBlockerReleased =
+                        mWakelockController.releaseProxDebounceSuspendBlocker();
+                if (proxDebounceSuspendBlockerReleased) {
+                    mPendingProximityDebounceTime = -1;
+                }
             }
         }
     }
@@ -2294,12 +2255,12 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
             mHandler.removeMessages(MSG_PROXIMITY_SENSOR_DEBOUNCED);
             if (positive) {
                 mPendingProximity = PROXIMITY_POSITIVE;
-                setPendingProximityDebounceTime(
-                        time + PROXIMITY_SENSOR_POSITIVE_DEBOUNCE_DELAY); // acquire wake lock
+                mPendingProximityDebounceTime = time + PROXIMITY_SENSOR_POSITIVE_DEBOUNCE_DELAY;
+                mWakelockController.acquireProxDebounceSuspendBlocker(); // acquire wake lock
             } else {
                 mPendingProximity = PROXIMITY_NEGATIVE;
-                setPendingProximityDebounceTime(
-                        time + PROXIMITY_SENSOR_NEGATIVE_DEBOUNCE_DELAY); // acquire wake lock
+                mPendingProximityDebounceTime = time + PROXIMITY_SENSOR_NEGATIVE_DEBOUNCE_DELAY;
+                mWakelockController.acquireProxDebounceSuspendBlocker(); // acquire wake lock
             }
 
             // Debounce the new sensor reading.
@@ -2321,7 +2282,13 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
                 // Sensor reading accepted.  Apply the change then release the wake lock.
                 mProximity = mPendingProximity;
                 updatePowerState();
-                clearPendingProximityDebounceTime(); // release wake lock (must be last)
+                // (must be last)
+                boolean proxDebounceSuspendBlockerReleased =
+                        mWakelockController.releaseProxDebounceSuspendBlocker();
+                if (proxDebounceSuspendBlockerReleased) {
+                    mPendingProximityDebounceTime = -1;
+                }
+
             } else {
                 // Need to wait a little longer.
                 // Debounce again later.  We continue holding a wake lock while waiting.
@@ -2331,25 +2298,10 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         }
     }
 
-    private void clearPendingProximityDebounceTime() {
-        if (mPendingProximityDebounceTime >= 0) {
-            mPendingProximityDebounceTime = -1;
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxDebounce);
-        }
-    }
-
-    private void setPendingProximityDebounceTime(long debounceTime) {
-        if (mPendingProximityDebounceTime < 0) {
-            mCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxDebounce);
-        }
-        mPendingProximityDebounceTime = debounceTime;
-    }
-
     private void sendOnStateChangedWithWakelock() {
-        if (!mOnStateChangedPending) {
-            mOnStateChangedPending = true;
-            mCallbacks.acquireSuspendBlocker(mSuspendBlockerIdOnStateChanged);
-            mHandler.post(mOnStateChangedRunnable);
+        boolean wakeLockAcquired = mWakelockController.acquireStateChangedSuspendBlocker();
+        if (wakeLockAcquired) {
+            mHandler.post(mWakelockController.getOnStateChangedRunnable());
         }
     }
 
@@ -2510,44 +2462,17 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         }
     }
 
-    private final Runnable mOnStateChangedRunnable = new Runnable() {
-        @Override
-        public void run() {
-            mOnStateChangedPending = false;
-            mCallbacks.onStateChanged();
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdOnStateChanged);
-        }
-    };
-
     private void sendOnProximityPositiveWithWakelock() {
-        mCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxPositive);
-        mHandler.post(mOnProximityPositiveRunnable);
-        mOnProximityPositiveMessages++;
+        mWakelockController.acquireProxPositiveSuspendBlocker();
+        mHandler.post(mWakelockController.getOnProximityPositiveRunnable());
     }
 
-    private final Runnable mOnProximityPositiveRunnable = new Runnable() {
-        @Override
-        public void run() {
-            mOnProximityPositiveMessages--;
-            mCallbacks.onProximityPositive();
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxPositive);
-        }
-    };
 
     private void sendOnProximityNegativeWithWakelock() {
-        mOnProximityNegativeMessages++;
-        mCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxNegative);
-        mHandler.post(mOnProximityNegativeRunnable);
+        mWakelockController.acquireProxNegativeSuspendBlocker();
+        mHandler.post(mWakelockController.getOnProximityNegativeRunnable());
     }
 
-    private final Runnable mOnProximityNegativeRunnable = new Runnable() {
-        @Override
-        public void run() {
-            mOnProximityNegativeMessages--;
-            mCallbacks.onProximityNegative();
-            mCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxNegative);
-        }
-    };
 
     @Override
     public void dump(final PrintWriter pw) {
@@ -2606,7 +2531,6 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         pw.println();
         pw.println("Display Power Controller Thread State:");
         pw.println("  mPowerRequest=" + mPowerRequest);
-        pw.println("  mUnfinishedBusiness=" + mUnfinishedBusiness);
         pw.println("  mWaitingForNegativeProximity=" + mWaitingForNegativeProximity);
         pw.println("  mProximitySensor=" + mProximitySensor);
         pw.println("  mProximitySensorEnabled=" + mProximitySensorEnabled);
@@ -2644,9 +2568,6 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         pw.println("  mReportedToPolicy="
                 + reportedToPolicyToString(mReportedScreenStateToPolicy));
         pw.println("  mIsRbcActive=" + mIsRbcActive);
-        pw.println("  mOnStateChangePending=" + mOnStateChangedPending);
-        pw.println("  mOnProximityPositiveMessages=" + mOnProximityPositiveMessages);
-        pw.println("  mOnProximityNegativeMessages=" + mOnProximityNegativeMessages);
 
         if (mScreenBrightnessRampAnimator != null) {
             pw.println("  mScreenBrightnessRampAnimator.isAnimating()="
@@ -2684,6 +2605,12 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
             mDisplayWhiteBalanceController.dump(pw);
             mDisplayWhiteBalanceSettings.dump(pw);
         }
+
+        pw.println();
+
+        if (mWakelockController != null) {
+            mWakelockController.dumpLocal(pw);
+        }
     }
 
     private static String proximityToString(int state) {
@@ -2989,28 +2916,6 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
         }
     }
 
-    @VisibleForTesting
-    String getSuspendBlockerUnfinishedBusinessId(int displayId) {
-        return "[" + displayId + "]unfinished business";
-    }
-
-    String getSuspendBlockerOnStateChangedId(int displayId) {
-        return "[" + displayId + "]on state changed";
-    }
-
-    String getSuspendBlockerProxPositiveId(int displayId) {
-        return "[" + displayId + "]prox positive";
-    }
-
-    String getSuspendBlockerProxNegativeId(int displayId) {
-        return "[" + displayId + "]prox negative";
-    }
-
-    @VisibleForTesting
-    String getSuspendBlockerProxDebounceId(int displayId) {
-        return "[" + displayId + "]prox debounce";
-    }
-
     /** Functional interface for providing time. */
     @VisibleForTesting
     interface Clock {
@@ -3036,6 +2941,11 @@ final class DisplayPowerController2 implements AutomaticBrightnessController.Cal
                 FloatProperty<DisplayPowerState> secondProperty) {
             return new DualRampAnimator(dps, firstProperty, secondProperty);
         }
+
+        WakelockController getWakelockController(int displayId,
+                DisplayPowerCallbacks displayPowerCallbacks) {
+            return new WakelockController(displayId, displayPowerCallbacks);
+        }
     }
 
     static class CachedBrightnessInfo {
diff --git a/services/core/java/com/android/server/display/WakelockController.java b/services/core/java/com/android/server/display/WakelockController.java
new file mode 100644
index 000000000000..cbf559fdb84f
--- /dev/null
+++ b/services/core/java/com/android/server/display/WakelockController.java
@@ -0,0 +1,302 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display;
+
+import android.hardware.display.DisplayManagerInternal;
+import android.util.Slog;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.io.PrintWriter;
+
+/**
+ * A utility class to acquire/release suspend blockers and manage appropriate states around it.
+ * It is also a channel to asynchronously update the PowerManagerService about the changes in the
+ * display states as needed.
+ */
+public final class WakelockController {
+    private static final boolean DEBUG = false;
+
+    // Asynchronous callbacks into the power manager service.
+    // Only invoked from the handler thread while no locks are held.
+    private final DisplayManagerInternal.DisplayPowerCallbacks mDisplayPowerCallbacks;
+
+    // Identifiers for suspend blocker acquisition requests
+    private final String mSuspendBlockerIdUnfinishedBusiness;
+    private final String mSuspendBlockerIdOnStateChanged;
+    private final String mSuspendBlockerIdProxPositive;
+    private final String mSuspendBlockerIdProxNegative;
+    private final String mSuspendBlockerIdProxDebounce;
+
+    // True if we have unfinished business and are holding a suspend-blocker.
+    private boolean mUnfinishedBusiness;
+
+    // True if we have have debounced the proximity change impact and are holding a suspend-blocker.
+    private boolean mHasProximityDebounced;
+
+    // The ID of the LogicalDisplay tied to this.
+    private final int mDisplayId;
+    private final String mTag;
+
+    // When true, it implies a wakelock is being held to guarantee the update happens before we
+    // collapse into suspend and so needs to be cleaned up if the thread is exiting.
+    // Should only be accessed on the Handler thread of the class managing the Display states
+    // (i.e. DisplayPowerController2).
+    private boolean mOnStateChangedPending;
+
+    // Count of positive proximity messages currently held. Used to keep track of how many
+    // suspend blocker acquisitions are pending when shutting down the DisplayPowerController2.
+    // Should only be accessed on the Handler thread of the class managing the Display states
+    // (i.e. DisplayPowerController2).
+    private int mOnProximityPositiveMessages;
+
+    // Count of negative proximity messages currently held. Used to keep track of how many
+    // suspend blocker acquisitions are pending when shutting down the DisplayPowerController2.
+    // Should only be accessed on the Handler thread of the class managing the Display states
+    // (i.e. DisplayPowerController2).
+    private int mOnProximityNegativeMessages;
+
+    /**
+     * The constructor of WakelockController. Manages the initialization of all the local entities
+     * needed for its appropriate functioning.
+     */
+    public WakelockController(int displayId,
+            DisplayManagerInternal.DisplayPowerCallbacks callbacks) {
+        mDisplayId = displayId;
+        mTag = "WakelockController[" + mDisplayId + "]";
+        mDisplayPowerCallbacks = callbacks;
+        mSuspendBlockerIdUnfinishedBusiness = "[" + displayId + "]unfinished business";
+        mSuspendBlockerIdOnStateChanged = "[" + displayId + "]on state changed";
+        mSuspendBlockerIdProxPositive = "[" + displayId + "]prox positive";
+        mSuspendBlockerIdProxNegative = "[" + displayId + "]prox negative";
+        mSuspendBlockerIdProxDebounce = "[" + displayId + "]prox debounce";
+    }
+
+    /**
+     * Acquires the state change wakelock and notifies the PowerManagerService about the changes.
+     */
+    public boolean acquireStateChangedSuspendBlocker() {
+        // Grab a wake lock if we have change of the display state
+        if (!mOnStateChangedPending) {
+            if (DEBUG) {
+                Slog.d(mTag, "State Changed...");
+            }
+            mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdOnStateChanged);
+            mOnStateChangedPending = true;
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * Releases the state change wakelock and notifies the PowerManagerService about the changes.
+     */
+    public void releaseStateChangedSuspendBlocker() {
+        if (mOnStateChangedPending) {
+            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdOnStateChanged);
+            mOnStateChangedPending = false;
+        }
+    }
+
+    /**
+     * Acquires the unfinished business wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    public void acquireUnfinishedBusinessSuspendBlocker() {
+        // Grab a wake lock if we have unfinished business.
+        if (!mUnfinishedBusiness) {
+            if (DEBUG) {
+                Slog.d(mTag, "Unfinished business...");
+            }
+            mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdUnfinishedBusiness);
+            mUnfinishedBusiness = true;
+        }
+    }
+
+    /**
+     * Releases the unfinished business wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    public void releaseUnfinishedBusinessSuspendBlocker() {
+        if (mUnfinishedBusiness) {
+            if (DEBUG) {
+                Slog.d(mTag, "Finished business...");
+            }
+            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdUnfinishedBusiness);
+            mUnfinishedBusiness = false;
+        }
+    }
+
+    /**
+     * Acquires the proximity positive wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    public void acquireProxPositiveSuspendBlocker() {
+        mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxPositive);
+        mOnProximityPositiveMessages++;
+    }
+
+    /**
+     * Releases the proximity positive wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    public void releaseProxPositiveSuspendBlocker() {
+        for (int i = 0; i < mOnProximityPositiveMessages; i++) {
+            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxPositive);
+        }
+        mOnProximityPositiveMessages = 0;
+    }
+
+    /**
+     * Acquires the proximity negative wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    public void acquireProxNegativeSuspendBlocker() {
+        mOnProximityNegativeMessages++;
+        mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxNegative);
+    }
+
+    /**
+     * Releases the proximity negative wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    public void releaseProxNegativeSuspendBlocker() {
+        for (int i = 0; i < mOnProximityNegativeMessages; i++) {
+            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxNegative);
+        }
+        mOnProximityNegativeMessages = 0;
+    }
+
+    /**
+     * Acquires the proximity debounce wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    public void acquireProxDebounceSuspendBlocker() {
+        if (!mHasProximityDebounced) {
+            mDisplayPowerCallbacks.acquireSuspendBlocker(mSuspendBlockerIdProxDebounce);
+        }
+        mHasProximityDebounced = true;
+    }
+
+    /**
+     * Releases the proximity debounce wakelock and notifies the PowerManagerService about the
+     * changes.
+     */
+    public boolean releaseProxDebounceSuspendBlocker() {
+        if (mHasProximityDebounced) {
+            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxDebounce);
+            mHasProximityDebounced = false;
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * Gets the Runnable to be executed when the proximity becomes positive.
+     */
+    public Runnable getOnProximityPositiveRunnable() {
+        return () -> {
+            mOnProximityPositiveMessages--;
+            mDisplayPowerCallbacks.onProximityPositive();
+            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxPositive);
+        };
+    }
+
+    /**
+     * Gets the Runnable to be executed when the display state changes
+     */
+    public Runnable getOnStateChangedRunnable() {
+        return () -> {
+            mOnStateChangedPending = false;
+            mDisplayPowerCallbacks.onStateChanged();
+            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdOnStateChanged);
+        };
+    }
+
+    /**
+     * Gets the Runnable to be executed when the proximity becomes negative.
+     */
+    public Runnable getOnProximityNegativeRunnable() {
+        return () -> {
+            mOnProximityNegativeMessages--;
+            mDisplayPowerCallbacks.onProximityNegative();
+            mDisplayPowerCallbacks.releaseSuspendBlocker(mSuspendBlockerIdProxNegative);
+        };
+    }
+
+    /**
+     * Dumps the current state of this
+     */
+    public void dumpLocal(PrintWriter pw) {
+        pw.println("WakelockController State:");
+        pw.println("  mDisplayId=" + mDisplayId);
+        pw.println("  mUnfinishedBusiness=" + hasUnfinishedBusiness());
+        pw.println("  mOnStateChangePending=" + isOnStateChangedPending());
+        pw.println("  mOnProximityPositiveMessages=" + getOnProximityPositiveMessages());
+        pw.println("  mOnProximityNegativeMessages=" + getOnProximityNegativeMessages());
+    }
+
+    @VisibleForTesting
+    String getSuspendBlockerUnfinishedBusinessId() {
+        return mSuspendBlockerIdUnfinishedBusiness;
+    }
+
+    @VisibleForTesting
+    String getSuspendBlockerOnStateChangedId() {
+        return mSuspendBlockerIdOnStateChanged;
+    }
+
+    @VisibleForTesting
+    String getSuspendBlockerProxPositiveId() {
+        return mSuspendBlockerIdProxPositive;
+    }
+
+    @VisibleForTesting
+    String getSuspendBlockerProxNegativeId() {
+        return mSuspendBlockerIdProxNegative;
+    }
+
+    @VisibleForTesting
+    String getSuspendBlockerProxDebounceId() {
+        return mSuspendBlockerIdProxDebounce;
+    }
+
+    @VisibleForTesting
+    boolean hasUnfinishedBusiness() {
+        return mUnfinishedBusiness;
+    }
+
+    @VisibleForTesting
+    boolean isOnStateChangedPending() {
+        return mOnStateChangedPending;
+    }
+
+    @VisibleForTesting
+    int getOnProximityPositiveMessages() {
+        return mOnProximityPositiveMessages;
+    }
+
+    @VisibleForTesting
+    int getOnProximityNegativeMessages() {
+        return mOnProximityNegativeMessages;
+    }
+
+    @VisibleForTesting
+    boolean hasProximitySensorDebounced() {
+        return mHasProximityDebounced;
+    }
+}
diff --git a/services/tests/mockingservicestests/src/com/android/server/display/DisplayPowerController2Test.java b/services/tests/mockingservicestests/src/com/android/server/display/DisplayPowerController2Test.java
index 1a5d496f2002..2df6823afab1 100644
--- a/services/tests/mockingservicestests/src/com/android/server/display/DisplayPowerController2Test.java
+++ b/services/tests/mockingservicestests/src/com/android/server/display/DisplayPowerController2Test.java
@@ -98,6 +98,8 @@ public final class DisplayPowerController2Test {
     private DisplayPowerState mDisplayPowerStateMock;
     @Mock
     private DualRampAnimator<DisplayPowerState> mDualRampAnimatorMock;
+    @Mock
+    private WakelockController mWakelockController;
 
     @Captor
     private ArgumentCaptor<SensorEventListener> mSensorEventListenerCaptor;
@@ -127,6 +129,12 @@ public final class DisplayPowerController2Test {
                     FloatProperty<DisplayPowerState> secondProperty) {
                 return mDualRampAnimatorMock;
             }
+
+            @Override
+            WakelockController getWakelockController(int displayId,
+                    DisplayPowerCallbacks displayPowerCallbacks) {
+                return mWakelockController;
+            }
         };
 
         addLocalServiceMock(WindowManagerPolicy.class, mWindowManagerPolicyMock);
@@ -169,19 +177,15 @@ public final class DisplayPowerController2Test {
         advanceTime(1);
 
         // two times, one for unfinished business and one for proximity
-        verify(mDisplayPowerCallbacksMock).acquireSuspendBlocker(
-                dpc.getSuspendBlockerUnfinishedBusinessId(DISPLAY_ID));
-        verify(mDisplayPowerCallbacksMock).acquireSuspendBlocker(
-                dpc.getSuspendBlockerProxDebounceId(DISPLAY_ID));
+        verify(mWakelockController).acquireUnfinishedBusinessSuspendBlocker();
+        verify(mWakelockController).acquireProxDebounceSuspendBlocker();
+
 
         dpc.stop();
         advanceTime(1);
-
         // two times, one for unfinished business and one for proximity
-        verify(mDisplayPowerCallbacksMock).releaseSuspendBlocker(
-                dpc.getSuspendBlockerUnfinishedBusinessId(DISPLAY_ID));
-        verify(mDisplayPowerCallbacksMock).releaseSuspendBlocker(
-                dpc.getSuspendBlockerProxDebounceId(DISPLAY_ID));
+        verify(mWakelockController).acquireUnfinishedBusinessSuspendBlocker();
+        verify(mWakelockController).acquireProxDebounceSuspendBlocker();
     }
 
     /**
diff --git a/services/tests/mockingservicestests/src/com/android/server/display/WakelockControllerTest.java b/services/tests/mockingservicestests/src/com/android/server/display/WakelockControllerTest.java
new file mode 100644
index 000000000000..288408c84a4f
--- /dev/null
+++ b/services/tests/mockingservicestests/src/com/android/server/display/WakelockControllerTest.java
@@ -0,0 +1,240 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.hardware.display.DisplayManagerInternal;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public final class WakelockControllerTest {
+    private static final int DISPLAY_ID = 1;
+
+    @Mock
+    private DisplayManagerInternal.DisplayPowerCallbacks mDisplayPowerCallbacks;
+
+    private WakelockController mWakelockController;
+
+    @Before
+    public void before() {
+        MockitoAnnotations.initMocks(this);
+        mWakelockController = new WakelockController(DISPLAY_ID, mDisplayPowerCallbacks);
+    }
+
+    @Test
+    public void validateSuspendBlockerIdsAreExpected() {
+        assertEquals(mWakelockController.getSuspendBlockerUnfinishedBusinessId(),
+                "[" + DISPLAY_ID + "]unfinished business");
+        assertEquals(mWakelockController.getSuspendBlockerOnStateChangedId(),
+                "[" + DISPLAY_ID + "]on state changed");
+        assertEquals(mWakelockController.getSuspendBlockerProxPositiveId(),
+                "[" + DISPLAY_ID + "]prox positive");
+        assertEquals(mWakelockController.getSuspendBlockerProxNegativeId(),
+                "[" + DISPLAY_ID + "]prox negative");
+        assertEquals(mWakelockController.getSuspendBlockerProxDebounceId(),
+                "[" + DISPLAY_ID + "]prox debounce");
+    }
+
+    @Test
+    public void acquireStateChangedSuspendBlockerAcquiresIfNotAcquired() {
+        // Acquire the suspend blocker
+        assertTrue(mWakelockController.acquireStateChangedSuspendBlocker());
+        assertTrue(mWakelockController.isOnStateChangedPending());
+
+        // Try to reacquire
+        assertFalse(mWakelockController.acquireStateChangedSuspendBlocker());
+        assertTrue(mWakelockController.isOnStateChangedPending());
+
+        // Verify acquire happened only once
+        verify(mDisplayPowerCallbacks, times(1))
+                .acquireSuspendBlocker(mWakelockController.getSuspendBlockerOnStateChangedId());
+
+        // Release
+        mWakelockController.releaseStateChangedSuspendBlocker();
+        assertFalse(mWakelockController.isOnStateChangedPending());
+
+        // Try to release again
+        mWakelockController.releaseStateChangedSuspendBlocker();
+
+        // Verify release happened only once
+        verify(mDisplayPowerCallbacks, times(1))
+                .releaseSuspendBlocker(mWakelockController.getSuspendBlockerOnStateChangedId());
+    }
+
+    @Test
+    public void acquireUnfinishedBusinessSuspendBlockerAcquiresIfNotAcquired() {
+        // Acquire the suspend blocker
+        mWakelockController.acquireUnfinishedBusinessSuspendBlocker();
+        assertTrue(mWakelockController.hasUnfinishedBusiness());
+
+        // Try to reacquire
+        mWakelockController.acquireUnfinishedBusinessSuspendBlocker();
+        assertTrue(mWakelockController.hasUnfinishedBusiness());
+
+        // Verify acquire happened only once
+        verify(mDisplayPowerCallbacks, times(1))
+                .acquireSuspendBlocker(mWakelockController.getSuspendBlockerUnfinishedBusinessId());
+
+        // Release the suspend blocker
+        mWakelockController.releaseUnfinishedBusinessSuspendBlocker();
+        assertFalse(mWakelockController.hasUnfinishedBusiness());
+
+        // Try to release again
+        mWakelockController.releaseUnfinishedBusinessSuspendBlocker();
+
+        // Verify release happened only once
+        verify(mDisplayPowerCallbacks, times(1))
+                .releaseSuspendBlocker(mWakelockController.getSuspendBlockerUnfinishedBusinessId());
+    }
+
+    @Test
+    public void acquireProxPositiveSuspendBlockerAcquiresIfNotAcquired() {
+        // Acquire the suspend blocker
+        mWakelockController.acquireProxPositiveSuspendBlocker();
+        assertEquals(mWakelockController.getOnProximityPositiveMessages(), 1);
+
+        // Try to reacquire
+        mWakelockController.acquireProxPositiveSuspendBlocker();
+        assertEquals(mWakelockController.getOnProximityPositiveMessages(), 2);
+
+        // Verify acquire happened only once
+        verify(mDisplayPowerCallbacks, times(2))
+                .acquireSuspendBlocker(mWakelockController.getSuspendBlockerProxPositiveId());
+
+        // Release the suspend blocker
+        mWakelockController.releaseProxPositiveSuspendBlocker();
+        assertEquals(mWakelockController.getOnProximityPositiveMessages(), 0);
+
+        // Verify all suspend blockers were released
+        verify(mDisplayPowerCallbacks, times(2))
+                .releaseSuspendBlocker(mWakelockController.getSuspendBlockerProxPositiveId());
+    }
+
+    @Test
+    public void acquireProxNegativeSuspendBlockerAcquiresIfNotAcquired() {
+        // Acquire the suspend blocker
+        mWakelockController.acquireProxNegativeSuspendBlocker();
+        assertEquals(mWakelockController.getOnProximityNegativeMessages(), 1);
+
+        // Try to reacquire
+        mWakelockController.acquireProxNegativeSuspendBlocker();
+        assertEquals(mWakelockController.getOnProximityNegativeMessages(), 2);
+
+        // Verify acquire happened only once
+        verify(mDisplayPowerCallbacks, times(2))
+                .acquireSuspendBlocker(mWakelockController.getSuspendBlockerProxNegativeId());
+
+        // Release the suspend blocker
+        mWakelockController.releaseProxNegativeSuspendBlocker();
+        assertEquals(mWakelockController.getOnProximityNegativeMessages(), 0);
+
+        // Verify all suspend blockers were released
+        verify(mDisplayPowerCallbacks, times(2))
+                .releaseSuspendBlocker(mWakelockController.getSuspendBlockerProxNegativeId());
+    }
+
+    @Test
+    public void acquireProxDebounceSuspendBlockerAcquiresIfNotAcquired() {
+        // Acquire the suspend blocker
+        mWakelockController.acquireProxDebounceSuspendBlocker();
+
+        // Try to reacquire
+        mWakelockController.acquireProxDebounceSuspendBlocker();
+        assertTrue(mWakelockController.hasProximitySensorDebounced());
+
+        // Verify acquire happened only once
+        verify(mDisplayPowerCallbacks, times(1))
+                .acquireSuspendBlocker(mWakelockController.getSuspendBlockerProxDebounceId());
+
+        // Release the suspend blocker
+        assertTrue(mWakelockController.releaseProxDebounceSuspendBlocker());
+
+        // Release again
+        assertFalse(mWakelockController.releaseProxDebounceSuspendBlocker());
+        assertFalse(mWakelockController.hasProximitySensorDebounced());
+
+        // Verify suspend blocker was released only once
+        verify(mDisplayPowerCallbacks, times(1))
+                .releaseSuspendBlocker(mWakelockController.getSuspendBlockerProxDebounceId());
+    }
+
+    @Test
+    public void proximityPositiveRunnableWorksAsExpected() {
+        // Acquire the suspend blocker twice
+        mWakelockController.acquireProxPositiveSuspendBlocker();
+        mWakelockController.acquireProxPositiveSuspendBlocker();
+
+        // Execute the runnable
+        Runnable proximityPositiveRunnable = mWakelockController.getOnProximityPositiveRunnable();
+        proximityPositiveRunnable.run();
+
+        // Validate one suspend blocker was released
+        assertEquals(mWakelockController.getOnProximityPositiveMessages(), 1);
+        verify(mDisplayPowerCallbacks).onProximityPositive();
+        verify(mDisplayPowerCallbacks).releaseSuspendBlocker(
+                mWakelockController.getSuspendBlockerProxPositiveId());
+    }
+
+    @Test
+    public void proximityNegativeRunnableWorksAsExpected() {
+        // Acquire the suspend blocker twice
+        mWakelockController.acquireProxNegativeSuspendBlocker();
+        mWakelockController.acquireProxNegativeSuspendBlocker();
+
+        // Execute the runnable
+        Runnable proximityNegativeRunnable = mWakelockController.getOnProximityNegativeRunnable();
+        proximityNegativeRunnable.run();
+
+        // Validate one suspend blocker was released
+        assertEquals(mWakelockController.getOnProximityNegativeMessages(), 1);
+        verify(mDisplayPowerCallbacks).onProximityNegative();
+        verify(mDisplayPowerCallbacks).releaseSuspendBlocker(
+                mWakelockController.getSuspendBlockerProxNegativeId());
+    }
+
+    @Test
+    public void onStateChangeRunnableWorksAsExpected() {
+        // Acquire the suspend blocker twice
+        mWakelockController.acquireStateChangedSuspendBlocker();
+
+        // Execute the runnable
+        Runnable stateChangeRunnable = mWakelockController.getOnStateChangedRunnable();
+        stateChangeRunnable.run();
+
+        // Validate one suspend blocker was released
+        assertEquals(mWakelockController.isOnStateChangedPending(), false);
+        verify(mDisplayPowerCallbacks).onStateChanged();
+        verify(mDisplayPowerCallbacks).releaseSuspendBlocker(
+                mWakelockController.getSuspendBlockerOnStateChangedId());
+    }
+
+
+}