From 55ef320bb379a7d41b001f48a12d24ddd78e56ef Mon Sep 17 00:00:00 2001
From: Yuncheol Heo <ycheo@google.com>
Date: Mon, 20 Sep 2021 17:04:14 -0700
Subject: [PATCH] Fix addListenerForType/removeListener in ShellTaskOranizer.

This CL fixes the following issues:
- addListenerForType for multi-types can invoke onTaskAppeared
  more than once.
- removeListener for mulit-types listener removes the first type only.

Bug: 200322591
Test: atest ShellTaskOrganizerTests
Change-Id: If53085ab229e3a1525b36aa3d964def0f228200c
---
 .../android/wm/shell/ShellTaskOrganizer.java  | 22 ++++++-----
 .../wm/shell/ShellTaskOrganizerTests.java     | 37 +++++++++++++++++++
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index d3265346036a0..d925a9218a77b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -233,14 +233,14 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
                             + " already exists");
                 }
                 mTaskListeners.put(listenerType, listener);
+            }
 
-                // Notify the listener of all existing tasks with the given type.
-                for (int i = mTasks.size() - 1; i >= 0; --i) {
-                    final TaskAppearedInfo data = mTasks.valueAt(i);
-                    final TaskListener taskListener = getTaskListener(data.getTaskInfo());
-                    if (taskListener != listener) continue;
-                    listener.onTaskAppeared(data.getTaskInfo(), data.getLeash());
-                }
+            // Notify the listener of all existing tasks with the given type.
+            for (int i = mTasks.size() - 1; i >= 0; --i) {
+                final TaskAppearedInfo data = mTasks.valueAt(i);
+                final TaskListener taskListener = getTaskListener(data.getTaskInfo());
+                if (taskListener != listener) continue;
+                listener.onTaskAppeared(data.getTaskInfo(), data.getLeash());
             }
         }
     }
@@ -266,8 +266,12 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
                 tasks.add(data);
             }
 
-            // Remove listener
-            mTaskListeners.removeAt(index);
+            // Remove listener, there can be the multiple occurrences, so search the whole list.
+            for (int i = mTaskListeners.size() - 1; i >= 0; --i) {
+                if (mTaskListeners.valueAt(i) == listener) {
+                    mTaskListeners.removeAt(i);
+                }
+            }
 
             // Associate tasks with new listeners if needed.
             for (int i = tasks.size() - 1; i >= 0; --i) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
index 6b74b620dad7b..d5acbbcf7d2c7 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
@@ -196,6 +196,43 @@ public class ShellTaskOrganizerTests {
         assertTrue(listener.appeared.contains(taskInfo));
     }
 
+    @Test
+    public void testAddListenerForMultipleTypes() {
+        RunningTaskInfo taskInfo1 = createTaskInfo(1, WINDOWING_MODE_FULLSCREEN);
+        mOrganizer.onTaskAppeared(taskInfo1, null);
+        RunningTaskInfo taskInfo2 = createTaskInfo(2, WINDOWING_MODE_MULTI_WINDOW);
+        mOrganizer.onTaskAppeared(taskInfo2, null);
+
+        TrackingTaskListener listener = new TrackingTaskListener();
+        mOrganizer.addListenerForType(listener,
+                TASK_LISTENER_TYPE_MULTI_WINDOW, TASK_LISTENER_TYPE_FULLSCREEN);
+
+        // onTaskAppeared event should be delivered once for each taskInfo.
+        assertTrue(listener.appeared.contains(taskInfo1));
+        assertTrue(listener.appeared.contains(taskInfo2));
+        assertEquals(2, listener.appeared.size());
+    }
+
+    @Test
+    public void testRemoveListenerForMultipleTypes() {
+        RunningTaskInfo taskInfo1 = createTaskInfo(1, WINDOWING_MODE_FULLSCREEN);
+        mOrganizer.onTaskAppeared(taskInfo1, null);
+        RunningTaskInfo taskInfo2 = createTaskInfo(2, WINDOWING_MODE_MULTI_WINDOW);
+        mOrganizer.onTaskAppeared(taskInfo2, null);
+
+        TrackingTaskListener listener = new TrackingTaskListener();
+        mOrganizer.addListenerForType(listener,
+                TASK_LISTENER_TYPE_MULTI_WINDOW, TASK_LISTENER_TYPE_FULLSCREEN);
+
+        mOrganizer.removeListener(listener);
+
+        // If listener is removed properly, onTaskInfoChanged event shouldn't be delivered.
+        mOrganizer.onTaskInfoChanged(taskInfo1);
+        assertTrue(listener.infoChanged.isEmpty());
+        mOrganizer.onTaskInfoChanged(taskInfo2);
+        assertTrue(listener.infoChanged.isEmpty());
+    }
+
     @Test
     public void testWindowingModeChange() {
         RunningTaskInfo taskInfo = createTaskInfo(1, WINDOWING_MODE_MULTI_WINDOW);