Remove all user data when a user is removed for Desktop Windowing.

Bug: 366397912
Flag: com.android.window.flags.enable_desktop_windowing_hsum
Test: DesktopPersistentRepositoryTest
Change-Id: I95546350c3459b04052f2bb4ad39e0d0d7ed3874

3 files changed, 91 insertions, 11 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopUserRepositories.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopUserRepositories.kt
index 7f3133e141ef..13576aa42737 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopUserRepositories.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopUserRepositories.kt
@@ -34,6 +34,7 @@ import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.sysui.UserChangeListener
 import java.io.PrintWriter
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 
 /** Manages per-user DesktopRepository instances. */
 class DesktopUserRepositories(
@@ -43,7 +44,7 @@ class DesktopUserRepositories(
     private val persistentRepository: DesktopPersistentRepository,
     private val repositoryInitializer: DesktopRepositoryInitializer,
     @ShellMainThread private val mainCoroutineScope: CoroutineScope,
-    userManager: UserManager,
+    private val userManager: UserManager,
 ) : UserChangeListener {
     private var userId: Int
     private var userIdToProfileIdsMap: MutableMap<Int, List<Int>> = mutableMapOf()
@@ -100,6 +101,9 @@ class DesktopUserRepositories(
     override fun onUserChanged(newUserId: Int, userContext: Context) {
         logD("onUserChanged previousUserId=%d, newUserId=%d", userId, newUserId)
         userId = newUserId
+        if (Flags.enableDesktopWindowingHsum()) {
+            sanitizeUsers()
+        }
     }
 
     override fun onUserProfilesChanged(profiles: MutableList<UserInfo>) {
@@ -110,10 +114,34 @@ class DesktopUserRepositories(
         }
     }
 
+    private fun sanitizeUsers() {
+        val aliveUserIds = userManager.getAliveUsers().map { it.id }
+        val usersToDelete = userIdToProfileIdsMap.keys.filterNot { it in aliveUserIds }
+
+        usersToDelete.forEach { uid ->
+            userIdToProfileIdsMap.remove(uid)
+            desktopRepoByUserId.remove(uid)
+        }
+        mainCoroutineScope.launch {
+            try {
+                persistentRepository.removeUsers(usersToDelete)
+            } catch (exception: Exception) {
+                logE(
+                    "An exception occurred while updating the persistent repository \n%s",
+                    exception.stackTrace,
+                )
+            }
+        }
+    }
+
     private fun logD(msg: String, vararg arguments: Any?) {
         ProtoLog.d(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
     }
 
+    private fun logE(msg: String, vararg arguments: Any?) {
+        ProtoLog.e(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
     companion object {
         private const val TAG = "DesktopUserRepositories"
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
index a6998e1179fa..9e41270c21f8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
@@ -141,6 +141,22 @@ class DesktopPersistentRepository(private val dataStore: DataStore<DesktopPersis
         }
     }
 
+    suspend fun removeUsers(uids: List<Int>) {
+        try {
+            dataStore.updateData { persistentRepositories: DesktopPersistentRepositories ->
+                val persistentRepositoriesBuilder = persistentRepositories.toBuilder()
+                uids.forEach { uid -> persistentRepositoriesBuilder.removeDesktopRepoByUser(uid) }
+                persistentRepositoriesBuilder.build()
+            }
+        } catch (exception: Exception) {
+            Log.e(
+                TAG,
+                "Error in removing user related data, data is stored in a file named $DESKTOP_REPOSITORIES_DATASTORE_FILE",
+                exception,
+            )
+        }
+    }
+
     private fun getDesktop(currentRepository: DesktopRepositoryState, desktopId: Int): Desktop =
         // If there are no desktops set up, create one on the default display
         currentRepository.getDesktopOrDefault(
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepositoryTest.kt
index eae206664021..5f92326b5e5e 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepositoryTest.kt
@@ -85,11 +85,11 @@ class DesktopPersistentRepositoryTest : ShellTestCase() {
             val desk = createDesktop(task)
             val repositoryState =
                 DesktopRepositoryState.newBuilder().putDesktop(DEFAULT_DESKTOP_ID, desk)
-            val DesktopPersistentRepositories =
+            val desktopPersistentRepositories =
                 DesktopPersistentRepositories.newBuilder()
                     .putDesktopRepoByUser(DEFAULT_USER_ID, repositoryState.build())
                     .build()
-            testDatastore.updateData { DesktopPersistentRepositories }
+            testDatastore.updateData { desktopPersistentRepositories }
 
             val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
 
@@ -102,8 +102,8 @@ class DesktopPersistentRepositoryTest : ShellTestCase() {
         runTest(StandardTestDispatcher()) {
             // Create a basic repository state
             val task = createDesktopTask(1)
-            val DesktopPersistentRepositories = createRepositoryWithOneDesk(task)
-            testDatastore.updateData { DesktopPersistentRepositories }
+            val desktopPersistentRepositories = createRepositoryWithOneDesk(task)
+            testDatastore.updateData { desktopPersistentRepositories }
             // Create a new state to be initialized
             val visibleTasks = ArraySet(listOf(1, 2))
             val minimizedTasks = ArraySet<Int>()
@@ -124,11 +124,46 @@ class DesktopPersistentRepositoryTest : ShellTestCase() {
     }
 
     @Test
+    fun removeUsers_removesUsersData() {
+        runTest(StandardTestDispatcher()) {
+            val task = createDesktopTask(1)
+            val desktopPersistentRepositories = createRepositoryWithOneDesk(task)
+            testDatastore.updateData { desktopPersistentRepositories }
+            // Create a new state to be initialized
+            val visibleTasks = ArraySet(listOf(1))
+            val minimizedTasks = ArraySet(listOf(1))
+            val freeformTasksInZOrder = ArrayList(listOf(1))
+            datastoreRepository.addOrUpdateDesktop(
+                visibleTasks = visibleTasks,
+                minimizedTasks = minimizedTasks,
+                freeformTasksInZOrder = freeformTasksInZOrder,
+                userId = DEFAULT_USER_ID,
+            )
+            datastoreRepository.addOrUpdateDesktop(
+                visibleTasks = visibleTasks,
+                minimizedTasks = minimizedTasks,
+                freeformTasksInZOrder = freeformTasksInZOrder,
+                userId = USER_ID_2,
+            )
+
+            datastoreRepository.removeUsers(mutableListOf(USER_ID_2))
+
+            val removedDesktopRepositoryState =
+                datastoreRepository.getDesktopRepositoryState(USER_ID_2)
+            assertThat(removedDesktopRepositoryState).isEqualTo(null)
+
+            val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
+            assertThat(actualDesktop?.tasksByTaskIdMap?.get(task.taskId)?.desktopTaskState)
+                .isEqualTo(DesktopTaskState.MINIMIZED)
+        }
+    }
+
+    @Test
     fun addOrUpdateTask_changeTaskStateToMinimize_taskStateIsMinimized() {
         runTest(StandardTestDispatcher()) {
             val task = createDesktopTask(1)
-            val DesktopPersistentRepositories = createRepositoryWithOneDesk(task)
-            testDatastore.updateData { DesktopPersistentRepositories }
+            val desktopPersistentRepositories = createRepositoryWithOneDesk(task)
+            testDatastore.updateData { desktopPersistentRepositories }
             // Create a new state to be initialized
             val visibleTasks = ArraySet(listOf(1))
             val minimizedTasks = ArraySet(listOf(1))
@@ -152,8 +187,8 @@ class DesktopPersistentRepositoryTest : ShellTestCase() {
     fun removeTask_previouslyAddedTaskIsRemoved() {
         runTest(StandardTestDispatcher()) {
             val task = createDesktopTask(1)
-            val DesktopPersistentRepositories = createRepositoryWithOneDesk(task)
-            testDatastore.updateData { DesktopPersistentRepositories }
+            val desktopPersistentRepositories = createRepositoryWithOneDesk(task)
+            testDatastore.updateData { desktopPersistentRepositories }
             // Create a new state to be initialized
             val visibleTasks = ArraySet<Int>()
             val minimizedTasks = ArraySet<Int>()
@@ -176,17 +211,18 @@ class DesktopPersistentRepositoryTest : ShellTestCase() {
     private companion object {
         const val DESKTOP_REPOSITORY_STATES_DATASTORE_TEST_FILE = "desktop_repo_test.pb"
         const val DEFAULT_USER_ID = 1000
+        const val USER_ID_2 = 2000
         const val DEFAULT_DESKTOP_ID = 0
 
         fun createRepositoryWithOneDesk(task: DesktopTask): DesktopPersistentRepositories {
             val desk = createDesktop(task)
             val repositoryState =
                 DesktopRepositoryState.newBuilder().putDesktop(DEFAULT_DESKTOP_ID, desk)
-            val DesktopPersistentRepositories =
+            val desktopPersistentRepositories =
                 DesktopPersistentRepositories.newBuilder()
                     .putDesktopRepoByUser(DEFAULT_USER_ID, repositoryState.build())
                     .build()
-            return DesktopPersistentRepositories
+            return desktopPersistentRepositories
         }
 
         fun createDesktop(task: DesktopTask): Desktop? =