From 63027a63d3b70f79a8b22e7d72cf49d7e0e30baa Mon Sep 17 00:00:00 2001
From: Dave Mankoff <mankoff@google.com>
Date: Thu, 21 May 2020 16:04:14 -0400
Subject: [PATCH] Ensure ProxCheck doesn't unregister ProxSensor

If the ProximitySensor passed into ProximityCheck is shared with other
clients, the ProximityCheck was causing the ProximitySensor to be
paused which, in turn, meant other clients would miss out on events.

Specifically, this manifested itself in AOD/Doze where one part of
Doze wanted the proximity sensor registered, but another part used
ProximityCheck with the same sensor, causing the sensor to be
paused. The screen would either not turn off when covered, or would
not turn back on when uncovered.

Fixes: 155365422
Test: atest SystemUITests and manual
Change-Id: I3c3faf79ba6d3915b0cfecec7861f6065ae689e8
---
 .../util/sensors/ProximitySensor.java         | 34 ++++++++++++-------
 .../util/sensors/ProximityCheckTest.java      | 26 ++++++++++++++
 2 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensor.java b/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensor.java
index 7445aa91eaed3..57f58c6a517ef 100644
--- a/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensor.java
+++ b/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensor.java
@@ -297,22 +297,15 @@ public class ProximitySensor implements ThresholdSensor {
         private final ProximitySensor mSensor;
         private final DelayableExecutor mDelayableExecutor;
         private List<Consumer<Boolean>> mCallbacks = new ArrayList<>();
+        private final ThresholdSensor.Listener mListener;
+        private final AtomicBoolean mRegistered = new AtomicBoolean();
 
         @Inject
         public ProximityCheck(ProximitySensor sensor, DelayableExecutor delayableExecutor) {
             mSensor = sensor;
             mSensor.setTag("prox_check");
             mDelayableExecutor = delayableExecutor;
-            mSensor.pause();
-            ThresholdSensor.Listener listener = proximityEvent -> {
-                mCallbacks.forEach(
-                        booleanConsumer ->
-                                booleanConsumer.accept(
-                                        proximityEvent == null ? null : proximityEvent.getBelow()));
-                mCallbacks.clear();
-                mSensor.pause();
-            };
-            mSensor.register(listener);
+            mListener = this::onProximityEvent;
         }
 
         /** Set a descriptive tag for the sensors registration. */
@@ -322,7 +315,7 @@ public class ProximitySensor implements ThresholdSensor {
 
         @Override
         public void run() {
-            mSensor.pause();
+            unregister();
             mSensor.alertListeners();
         }
 
@@ -334,11 +327,26 @@ public class ProximitySensor implements ThresholdSensor {
                 callback.accept(null);
             }
             mCallbacks.add(callback);
-            if (!mSensor.isRegistered()) {
-                mSensor.resume();
+            if (!mRegistered.getAndSet(true)) {
+                mSensor.register(mListener);
                 mDelayableExecutor.executeDelayed(this, timeoutMs);
             }
         }
+
+        private void unregister() {
+            mSensor.unregister(mListener);
+            mRegistered.set(false);
+        }
+
+        private void onProximityEvent(ThresholdSensorEvent proximityEvent) {
+            mCallbacks.forEach(
+                    booleanConsumer ->
+                            booleanConsumer.accept(
+                                    proximityEvent == null ? null : proximityEvent.getBelow()));
+            mCallbacks.clear();
+            unregister();
+            mRegistered.set(false);
+        }
     }
 
     private void logDebug(String msg) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/sensors/ProximityCheckTest.java b/packages/SystemUI/tests/src/com/android/systemui/util/sensors/ProximityCheckTest.java
index d83083bd12033..8ba11dae2b5c2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/sensors/ProximityCheckTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/sensors/ProximityCheckTest.java
@@ -82,6 +82,32 @@ public class ProximityCheckTest extends SysuiTestCase {
         assertFalse(mFakeProximitySensor.isRegistered());
     }
 
+    @Test
+    public void testProxDoesntCancelOthers() {
+        assertFalse(mFakeProximitySensor.isRegistered());
+        // We don't need our "other" listener to do anything. Just ensure our sensor is registered.
+        ThresholdSensor.Listener emptyListener = event -> { };
+        mFakeProximitySensor.register(emptyListener);
+        assertTrue(mFakeProximitySensor.isRegistered());
+
+        // Now run a basic check. This is just like testCheck()
+        mProximityCheck.check(100, mTestableCallback);
+
+        assertNull(mTestableCallback.mLastResult);
+
+        mFakeProximitySensor.setLastEvent(new ProximitySensor.ThresholdSensorEvent(true, 0));
+        mFakeProximitySensor.alertListeners();
+
+        assertTrue(mTestableCallback.mLastResult);
+
+        // We should still be registered, since we have another listener.
+        assertTrue(mFakeProximitySensor.isRegistered());
+
+        mFakeProximitySensor.unregister(emptyListener);
+        assertFalse(mFakeProximitySensor.isRegistered());
+
+    }
+
     private static class TestableCallback implements Consumer<Boolean> {
         Boolean mLastResult;
         @Override