Move active task tracking to a repository

We were tracking active freeform tasks in the RecentTasksController.
This causes some issues when we want to share that information to other
controllers like DesktopModeController.
Move tracking of active tasks to a separate repository that controllers
can query and listen for updates.

Bug: 244348395
Test: atest DesktopModeTaskRepositoryTest RecentTasksControllerTest
Change-Id: I7a460020f6b2133225e292a04f29ae1dd27e964b

7 files changed, 247 insertions, 39 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index c39602032170..7c3c14e2210c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -55,6 +55,8 @@ import com.android.wm.shell.common.annotations.ShellBackgroundThread;
 import com.android.wm.shell.common.annotations.ShellMainThread;
 import com.android.wm.shell.common.annotations.ShellSplashscreenThread;
 import com.android.wm.shell.compatui.CompatUIController;
+import com.android.wm.shell.desktopmode.DesktopMode;
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.displayareahelper.DisplayAreaHelper;
 import com.android.wm.shell.displayareahelper.DisplayAreaHelperController;
 import com.android.wm.shell.draganddrop.DragAndDropController;
@@ -477,11 +479,12 @@ public abstract class WMShellBaseModule {
             ShellInit shellInit,
             ShellCommandHandler shellCommandHandler,
             TaskStackListenerImpl taskStackListener,
+            Optional<DesktopModeTaskRepository> desktopModeTaskRepository,
             @ShellMainThread ShellExecutor mainExecutor
     ) {
         return Optional.ofNullable(
                 RecentTasksController.create(context, shellInit, shellCommandHandler,
-                        taskStackListener, mainExecutor));
+                        taskStackListener, desktopModeTaskRepository, mainExecutor));
     }
 
     //
@@ -666,6 +669,24 @@ public abstract class WMShellBaseModule {
     }
 
     //
+    // Desktop mode (optional feature)
+    //
+
+    @BindsOptionalOf
+    @DynamicOverride
+    abstract DesktopModeTaskRepository optionalDesktopModeTaskRepository();
+
+    @WMSingleton
+    @Provides
+    static Optional<DesktopModeTaskRepository> providesDesktopModeTaskRepository(
+            @DynamicOverride Optional<DesktopModeTaskRepository> desktopModeTaskRepository) {
+        if (DesktopMode.IS_SUPPORTED) {
+            return desktopModeTaskRepository;
+        }
+        return Optional.empty();
+    }
+
+    //
     // Misc
     //
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 18ce3642335d..639087ed9732 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -51,6 +51,7 @@ import com.android.wm.shell.common.annotations.ShellBackgroundThread;
 import com.android.wm.shell.common.annotations.ShellMainThread;
 import com.android.wm.shell.desktopmode.DesktopMode;
 import com.android.wm.shell.desktopmode.DesktopModeController;
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.freeform.FreeformComponents;
 import com.android.wm.shell.freeform.FreeformTaskListener;
@@ -220,14 +221,14 @@ public abstract class WMShellModule {
             Context context,
             ShellInit shellInit,
             ShellTaskOrganizer shellTaskOrganizer,
-            Optional<RecentTasksController> recentTasksController,
+            Optional<DesktopModeTaskRepository> desktopModeTaskRepository,
             WindowDecorViewModel<?> windowDecorViewModel) {
         // TODO(b/238217847): Temporarily add this check here until we can remove the dynamic
         //                    override for this controller from the base module
         ShellInit init = FreeformComponents.isFreeformEnabled(context)
                 ? shellInit
                 : null;
-        return new FreeformTaskListener<>(init, shellTaskOrganizer, recentTasksController,
+        return new FreeformTaskListener<>(init, shellTaskOrganizer, desktopModeTaskRepository,
                 windowDecorViewModel);
     }
 
@@ -610,6 +611,13 @@ public abstract class WMShellModule {
         }
     }
 
+    @WMSingleton
+    @Provides
+    @DynamicOverride
+    static DesktopModeTaskRepository provideDesktopModeTaskRepository() {
+        return new DesktopModeTaskRepository();
+    }
+
     //
     // Misc
     //
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
new file mode 100644
index 000000000000..988601c0e8a8
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
@@ -0,0 +1,89 @@
+/*
+ * 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.wm.shell.desktopmode
+
+import android.util.ArraySet
+
+/**
+ * Keeps track of task data related to desktop mode.
+ */
+class DesktopModeTaskRepository {
+
+    /**
+     * Set of task ids that are marked as active in desktop mode.
+     * Active tasks in desktop mode are freeform tasks that are visible or have been visible after
+     * desktop mode was activated.
+     * Task gets removed from this list when it vanishes. Or when desktop mode is turned off.
+     */
+    private val activeTasks = ArraySet<Int>()
+    private val listeners = ArraySet<Listener>()
+
+    /**
+     * Add a [Listener] to be notified of updates to the repository.
+     */
+    fun addListener(listener: Listener) {
+        listeners.add(listener)
+    }
+
+    /**
+     * Remove a previously registered [Listener]
+     */
+    fun removeListener(listener: Listener) {
+        listeners.remove(listener)
+    }
+
+    /**
+     * Mark a task with given [taskId] as active.
+     */
+    fun addActiveTask(taskId: Int) {
+        val added = activeTasks.add(taskId)
+        if (added) {
+            listeners.onEach { it.onActiveTasksChanged() }
+        }
+    }
+
+    /**
+     * Remove task with given [taskId] from active tasks.
+     */
+    fun removeActiveTask(taskId: Int) {
+        val removed = activeTasks.remove(taskId)
+        if (removed) {
+            listeners.onEach { it.onActiveTasksChanged() }
+        }
+    }
+
+    /**
+     * Check if a task with the given [taskId] was marked as an active task
+     */
+    fun isActiveTask(taskId: Int): Boolean {
+        return activeTasks.contains(taskId)
+    }
+
+    /**
+     * Get a set of the active tasks
+     */
+    fun getActiveTasks(): ArraySet<Int> {
+        return ArraySet(activeTasks)
+    }
+
+    /**
+     * Defines interface for classes that can listen to changes in repository state.
+     */
+    interface Listener {
+        fun onActiveTasksChanged()
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
index 1baac718ee95..ac95d4bc63d1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
@@ -29,8 +29,8 @@ import androidx.annotation.Nullable;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.desktopmode.DesktopMode;
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
-import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.windowdecor.WindowDecorViewModel;
@@ -49,7 +49,7 @@ public class FreeformTaskListener<T extends AutoCloseable>
     private static final String TAG = "FreeformTaskListener";
 
     private final ShellTaskOrganizer mShellTaskOrganizer;
-    private final Optional<RecentTasksController> mRecentTasksOptional;
+    private final Optional<DesktopModeTaskRepository> mDesktopModeTaskRepository;
     private final WindowDecorViewModel<T> mWindowDecorationViewModel;
 
     private final SparseArray<State<T>> mTasks = new SparseArray<>();
@@ -64,11 +64,11 @@ public class FreeformTaskListener<T extends AutoCloseable>
     public FreeformTaskListener(
             ShellInit shellInit,
             ShellTaskOrganizer shellTaskOrganizer,
-            Optional<RecentTasksController> recentTasksController,
+            Optional<DesktopModeTaskRepository> desktopModeTaskRepository,
             WindowDecorViewModel<T> windowDecorationViewModel) {
         mShellTaskOrganizer = shellTaskOrganizer;
         mWindowDecorationViewModel = windowDecorationViewModel;
-        mRecentTasksOptional = recentTasksController;
+        mDesktopModeTaskRepository = desktopModeTaskRepository;
         if (shellInit != null) {
             shellInit.addInitCallback(this::onInit, this);
         }
@@ -93,7 +93,7 @@ public class FreeformTaskListener<T extends AutoCloseable>
         if (DesktopMode.IS_SUPPORTED && taskInfo.isVisible) {
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                     "Adding active freeform task: #%d", taskInfo.taskId);
-            mRecentTasksOptional.ifPresent(rt -> rt.addActiveFreeformTask(taskInfo.taskId));
+            mDesktopModeTaskRepository.ifPresent(it -> it.addActiveTask(taskInfo.taskId));
         }
     }
 
@@ -126,7 +126,7 @@ public class FreeformTaskListener<T extends AutoCloseable>
         if (DesktopMode.IS_SUPPORTED) {
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                     "Removing active freeform task: #%d", taskInfo.taskId);
-            mRecentTasksOptional.ifPresent(rt -> rt.removeActiveFreeformTask(taskInfo.taskId));
+            mDesktopModeTaskRepository.ifPresent(it -> it.removeActiveTask(taskInfo.taskId));
         }
 
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
@@ -154,7 +154,7 @@ public class FreeformTaskListener<T extends AutoCloseable>
             if (taskInfo.isVisible) {
                 ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                         "Adding active freeform task: #%d", taskInfo.taskId);
-                mRecentTasksOptional.ifPresent(rt -> rt.addActiveFreeformTask(taskInfo.taskId));
+                mDesktopModeTaskRepository.ifPresent(it -> it.addActiveTask(taskInfo.taskId));
             }
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
index 27bc1a189086..ff4b2edb7310 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
@@ -44,6 +44,7 @@ import com.android.wm.shell.common.TaskStackListenerImpl;
 import com.android.wm.shell.common.annotations.ExternalThread;
 import com.android.wm.shell.common.annotations.ShellMainThread;
 import com.android.wm.shell.desktopmode.DesktopMode;
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellInit;
@@ -53,19 +54,20 @@ import com.android.wm.shell.util.SplitBounds;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 /**
  * Manages the recent task list from the system, caching it as necessary.
  */
 public class RecentTasksController implements TaskStackListenerCallback,
-        RemoteCallable<RecentTasksController> {
+        RemoteCallable<RecentTasksController>, DesktopModeTaskRepository.Listener {
     private static final String TAG = RecentTasksController.class.getSimpleName();
 
     private final Context mContext;
     private final ShellCommandHandler mShellCommandHandler;
+    private final Optional<DesktopModeTaskRepository> mDesktopModeTaskRepository;
     private final ShellExecutor mMainExecutor;
     private final TaskStackListenerImpl mTaskStackListener;
     private final RecentTasks mImpl = new RecentTasksImpl();
@@ -84,15 +86,6 @@ public class RecentTasksController implements TaskStackListenerCallback,
     private final Map<Integer, SplitBounds> mTaskSplitBoundsMap = new HashMap<>();
 
     /**
-     * Set of taskId's that have been launched in freeform mode.
-     * This includes tasks that are currently running, visible and in freeform mode. And also
-     * includes tasks that are running in the background, are no longer visible, but at some point
-     * were visible to the user.
-     * This is used to decide which freeform apps belong to the user's desktop.
-     */
-    private final HashSet<Integer> mActiveFreeformTasks = new HashSet<>();
-
-    /**
      * Creates {@link RecentTasksController}, returns {@code null} if the feature is not
      * supported.
      */
@@ -102,24 +95,27 @@ public class RecentTasksController implements TaskStackListenerCallback,
             ShellInit shellInit,
             ShellCommandHandler shellCommandHandler,
             TaskStackListenerImpl taskStackListener,
+            Optional<DesktopModeTaskRepository> desktopModeTaskRepository,
             @ShellMainThread ShellExecutor mainExecutor
     ) {
         if (!context.getResources().getBoolean(com.android.internal.R.bool.config_hasRecents)) {
             return null;
         }
         return new RecentTasksController(context, shellInit, shellCommandHandler, taskStackListener,
-                mainExecutor);
+                desktopModeTaskRepository, mainExecutor);
     }
 
     RecentTasksController(Context context,
             ShellInit shellInit,
             ShellCommandHandler shellCommandHandler,
             TaskStackListenerImpl taskStackListener,
+            Optional<DesktopModeTaskRepository> desktopModeTaskRepository,
             ShellExecutor mainExecutor) {
         mContext = context;
         mShellCommandHandler = shellCommandHandler;
         mIsDesktopMode = mContext.getPackageManager().hasSystemFeature(FEATURE_PC);
         mTaskStackListener = taskStackListener;
+        mDesktopModeTaskRepository = desktopModeTaskRepository;
         mMainExecutor = mainExecutor;
         shellInit.addInitCallback(this::onInit, this);
     }
@@ -131,6 +127,7 @@ public class RecentTasksController implements TaskStackListenerCallback,
     private void onInit() {
         mShellCommandHandler.addDumpCallback(this::dump, this);
         mTaskStackListener.addListener(this);
+        mDesktopModeTaskRepository.ifPresent(it -> it.addListener(this));
     }
 
     /**
@@ -217,19 +214,8 @@ public class RecentTasksController implements TaskStackListenerCallback,
         notifyRecentTasksChanged();
     }
 
-    /**
-     * Mark a task with given {@code taskId} as active in freeform
-     */
-    public void addActiveFreeformTask(int taskId) {
-        mActiveFreeformTasks.add(taskId);
-        notifyRecentTasksChanged();
-    }
-
-    /**
-     * Remove task with given {@code taskId} from active freeform tasks
-     */
-    public void removeActiveFreeformTask(int taskId) {
-        mActiveFreeformTasks.remove(taskId);
+    @Override
+    public void onActiveTasksChanged() {
         notifyRecentTasksChanged();
     }
 
@@ -312,7 +298,8 @@ public class RecentTasksController implements TaskStackListenerCallback,
                 continue;
             }
 
-            if (desktopModeActive && mActiveFreeformTasks.contains(taskInfo.taskId)) {
+            if (desktopModeActive && mDesktopModeTaskRepository.isPresent()
+                    && mDesktopModeTaskRepository.get().isActiveTask(taskInfo.taskId)) {
                 // Freeform tasks will be added as a separate entry
                 freeformTasks.add(taskInfo);
                 continue;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
new file mode 100644
index 000000000000..9b28d11f6a9d
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
@@ -0,0 +1,99 @@
+/*
+ * 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.wm.shell.desktopmode
+
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.wm.shell.ShellTestCase
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class DesktopModeTaskRepositoryTest : ShellTestCase() {
+
+    private lateinit var repo: DesktopModeTaskRepository
+
+    @Before
+    fun setUp() {
+        repo = DesktopModeTaskRepository()
+    }
+
+    @Test
+    fun addActiveTask_listenerNotifiedAndTaskIsActive() {
+        val listener = TestListener()
+        repo.addListener(listener)
+
+        repo.addActiveTask(1)
+        assertThat(listener.activeTaskChangedCalls).isEqualTo(1)
+        assertThat(repo.isActiveTask(1)).isTrue()
+    }
+
+    @Test
+    fun addActiveTask_sameTaskDoesNotNotify() {
+        val listener = TestListener()
+        repo.addListener(listener)
+
+        repo.addActiveTask(1)
+        repo.addActiveTask(1)
+        assertThat(listener.activeTaskChangedCalls).isEqualTo(1)
+    }
+
+    @Test
+    fun addActiveTask_multipleTasksAddedNotifiesForEach() {
+        val listener = TestListener()
+        repo.addListener(listener)
+
+        repo.addActiveTask(1)
+        repo.addActiveTask(2)
+        assertThat(listener.activeTaskChangedCalls).isEqualTo(2)
+    }
+
+    @Test
+    fun removeActiveTask_listenerNotifiedAndTaskNotActive() {
+        val listener = TestListener()
+        repo.addListener(listener)
+
+        repo.addActiveTask(1)
+        repo.removeActiveTask(1)
+        // Notify once for add and once for remove
+        assertThat(listener.activeTaskChangedCalls).isEqualTo(2)
+        assertThat(repo.isActiveTask(1)).isFalse()
+    }
+
+    @Test
+    fun removeActiveTask_removeNotExistingTaskDoesNotNotify() {
+        val listener = TestListener()
+        repo.addListener(listener)
+        repo.removeActiveTask(99)
+        assertThat(listener.activeTaskChangedCalls).isEqualTo(0)
+    }
+
+    @Test
+    fun isActiveTask_notExistingTaskReturnsFalse() {
+        assertThat(repo.isActiveTask(99)).isFalse()
+    }
+
+    class TestListener : DesktopModeTaskRepository.Listener {
+        var activeTaskChangedCalls = 0
+        override fun onActiveTasksChanged() {
+            activeTaskChangedCalls++
+        }
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
index e9a1e2523a86..cadfeb0de312 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
@@ -55,6 +55,7 @@ import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TaskStackListenerImpl;
 import com.android.wm.shell.desktopmode.DesktopMode;
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.util.GroupedRecentTaskInfo;
@@ -82,6 +83,8 @@ public class RecentTasksControllerTest extends ShellTestCase {
     private TaskStackListenerImpl mTaskStackListener;
     @Mock
     private ShellCommandHandler mShellCommandHandler;
+    @Mock
+    private DesktopModeTaskRepository mDesktopModeTaskRepository;
 
     private ShellTaskOrganizer mShellTaskOrganizer;
     private RecentTasksController mRecentTasksController;
@@ -94,7 +97,8 @@ public class RecentTasksControllerTest extends ShellTestCase {
         when(mContext.getPackageManager()).thenReturn(mock(PackageManager.class));
         mShellInit = spy(new ShellInit(mMainExecutor));
         mRecentTasksController = spy(new RecentTasksController(mContext, mShellInit,
-                mShellCommandHandler, mTaskStackListener, mMainExecutor));
+                mShellCommandHandler, mTaskStackListener, Optional.of(mDesktopModeTaskRepository),
+                mMainExecutor));
         mShellTaskOrganizer = new ShellTaskOrganizer(mShellInit, mShellCommandHandler,
                 null /* sizeCompatUI */, Optional.empty(), Optional.of(mRecentTasksController),
                 mMainExecutor);
@@ -195,8 +199,8 @@ public class RecentTasksControllerTest extends ShellTestCase {
         ActivityManager.RecentTaskInfo t4 = makeTaskInfo(4);
         setRawList(t1, t2, t3, t4);
 
-        mRecentTasksController.addActiveFreeformTask(1);
-        mRecentTasksController.addActiveFreeformTask(3);
+        when(mDesktopModeTaskRepository.isActiveTask(1)).thenReturn(true);
+        when(mDesktopModeTaskRepository.isActiveTask(3)).thenReturn(true);
 
         ArrayList<GroupedRecentTaskInfo> recentTasks = mRecentTasksController.getRecentTasks(
                 MAX_VALUE, RECENT_IGNORE_UNAVAILABLE, 0);