2 files changed, 50 insertions, 14 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 708b5a7a45f7..6794a2a3f72f 100644
--- a/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensor.java
+++ b/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensor.java
@@ -222,7 +222,9 @@ public class ProximitySensor {
         if (mAlerting.getAndSet(true)) {
             return;
         }
-        mListeners.forEach(proximitySensorListener ->
+
+        List<ProximitySensorListener> listeners = new ArrayList<>(mListeners);
+        listeners.forEach(proximitySensorListener ->
                 proximitySensorListener.onSensorEvent(mLastEvent));
         mAlerting.set(false);
     }
@@ -247,22 +249,15 @@ public class ProximitySensor {
         private final ProximitySensor mSensor;
         private final DelayableExecutor mDelayableExecutor;
         private List<Consumer<Boolean>> mCallbacks = new ArrayList<>();
+        private final ProximitySensor.ProximitySensorListener mListener;
+        private final AtomicBoolean mRegistered = new AtomicBoolean();
 
         @Inject
         public ProximityCheck(ProximitySensor sensor, DelayableExecutor delayableExecutor) {
             mSensor = sensor;
             mSensor.setTag("prox_check");
             mDelayableExecutor = delayableExecutor;
-            mSensor.pause();
-            ProximitySensorListener listener = proximityEvent -> {
-                mCallbacks.forEach(
-                        booleanConsumer ->
-                                booleanConsumer.accept(
-                                        proximityEvent == null ? null : proximityEvent.getNear()));
-                mCallbacks.clear();
-                mSensor.pause();
-            };
-            mSensor.register(listener);
+            mListener = this::onProximityEvent;
         }
 
         /** Set a descriptive tag for the sensors registration. */
@@ -272,7 +267,7 @@ public class ProximitySensor {
 
         @Override
         public void run() {
-            mSensor.pause();
+            unregister();
             mSensor.alertListeners();
         }
 
@@ -284,11 +279,26 @@ public class ProximitySensor {
                 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(ProximityEvent proximityEvent) {
+            mCallbacks.forEach(
+                    booleanConsumer ->
+                            booleanConsumer.accept(
+                                    proximityEvent == null ? null : proximityEvent.getNear()));
+            mCallbacks.clear();
+            unregister();
+            mRegistered.set(false);
+        }
     }
 
     /** Implement to be notified of ProximityEvents. */
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 7221095a1eaf..dae6b284ff15 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
@@ -81,6 +81,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.
+        ProximitySensor.ProximitySensorListener 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.ProximityEvent(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