From 0a3f4b80e575ca9ff9da02f006664706c9578c09 Mon Sep 17 00:00:00 2001
From: Bryce Lee <brycelee@google.com>
Date: Thu, 19 Jan 2023 16:15:00 -0800
Subject: [PATCH] Ensure TouchSessions are always removed.

Gesture events normally drive TouchSession cleanup with up/cancel
events triggering the removal of any active sessions. However, it
is possible for the touch monitor to be stopped before this point.
This changelist ensure that the sessions are properly cleaned up
when monitoring stops before the expected touch events.

Test: atest DreamOverlayTouchMonitorTest#testOnRemovedCallbackOnStopMonitor
Fixes: 245441667
Change-Id: Id6817e726058b0a2738cbb7c163d2efa5eb3875e
---
 .../touch/DreamOverlayTouchMonitor.java       |  9 +++++++
 .../touch/DreamOverlayTouchMonitorTest.java   | 26 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/packages/SystemUI/src/com/android/systemui/dreams/touch/DreamOverlayTouchMonitor.java b/packages/SystemUI/src/com/android/systemui/dreams/touch/DreamOverlayTouchMonitor.java
index 695b59ac45ad8..b8b459e1c68cb 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/touch/DreamOverlayTouchMonitor.java
+++ b/packages/SystemUI/src/com/android/systemui/dreams/touch/DreamOverlayTouchMonitor.java
@@ -226,6 +226,15 @@ public class DreamOverlayTouchMonitor {
             return;
         }
 
+        // When we stop monitoring touches, we must ensure that all active touch sessions and
+        // descendants informed of the removal so any cleanup for active tracking can proceed.
+        mExecutor.execute(() -> mActiveTouchSessions.forEach(touchSession -> {
+            while (touchSession != null) {
+                touchSession.onRemoved();
+                touchSession = touchSession.getPredecessor();
+            }
+        }));
+
         mCurrentInputSession.dispose();
         mCurrentInputSession = null;
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/dreams/touch/DreamOverlayTouchMonitorTest.java b/packages/SystemUI/tests/src/com/android/systemui/dreams/touch/DreamOverlayTouchMonitorTest.java
index a807407f170d1..178b9cc727269 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/dreams/touch/DreamOverlayTouchMonitorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/dreams/touch/DreamOverlayTouchMonitorTest.java
@@ -424,6 +424,32 @@ public class DreamOverlayTouchMonitorTest extends SysuiTestCase {
         verify(gestureListener2).onDown(eq(followupEvent));
     }
 
+    @Test
+    public void testOnRemovedCallbackOnStopMonitoring() {
+        final DreamTouchHandler touchHandler = Mockito.mock(DreamTouchHandler.class);
+        final DreamTouchHandler.TouchSession.Callback callback =
+                Mockito.mock(DreamTouchHandler.TouchSession.Callback.class);
+
+        final Environment environment = new Environment(Stream.of(touchHandler)
+                .collect(Collectors.toCollection(HashSet::new)));
+
+        final InputEvent initialEvent = Mockito.mock(InputEvent.class);
+        environment.publishInputEvent(initialEvent);
+
+        final DreamTouchHandler.TouchSession session = captureSession(touchHandler);
+        session.registerCallback(callback);
+
+        environment.executeAll();
+
+        environment.updateLifecycle(observerOwnerPair -> {
+            observerOwnerPair.first.onPause(observerOwnerPair.second);
+        });
+
+        environment.executeAll();
+
+        verify(callback).onRemoved();
+    }
+
     public GestureDetector.OnGestureListener registerGestureListener(DreamTouchHandler handler) {
         final GestureDetector.OnGestureListener gestureListener = Mockito.mock(
                 GestureDetector.OnGestureListener.class);