Donate to e Foundation | Murena handsets with /e/OS | Own a part of Murena! Learn more

Commit af3a1c4b authored by Orhan Uysal's avatar Orhan Uysal
Browse files

Persist all of the repo state. 

When there is a change in the repository, persist all of the desks
rather than just the desk that has had a change. This is needed because
back to back partial updates to the DataStore can create issues when
there are multiple removals.

Bug: 425916459
Test: atest DesktopRepositoryTest
Test: atest DesktopPersistentRepositoryTest
Flag: com.android.window.flags.repository_based_persistence

Change-Id: I57c0129d0ef65d357235beea6a9d77c6ec945d97
parent 788876db

libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/data/DesktopRepository.kt

+81 −18
Original line number Diff line number Diff line
@@ -284,6 +284,9 @@ class DesktopRepository( 
    /** Returns all the ids of all desks in all displays. */

    fun getAllDeskIds(): Set<Int> = desktopData.desksSequence().map { desk -> desk.deskId }.toSet()



    /** Returns all the desks in all displays. */

    @VisibleForTesting fun getAllDesks(): List<Desk> = desktopData.desksSequence().toList()



    /** Returns the id of the default desk in the given display. */

    fun getDefaultDeskId(displayId: Int): Int? = getDefaultDesk(displayId)?.deskId
@@ -361,10 +364,14 @@ class DesktopRepository( 
        )

        val desk = checkNotNull(desktopData.getDesk(deskId)) { "Did not find desk: $deskId" }

        desk.leftTiledTaskId = taskId

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {

            if (DesktopExperienceFlags.REPOSITORY_BASED_PERSISTENCE.isTrue) {

                updatePersistentRepository(displayId)

            } else {

                updatePersistentRepositoryForDesk(deskId)

            }

        }

    }



    fun addRightTiledTaskToDesk(displayId: Int, taskId: Int, deskId: Int) {

        logD(
@@ -375,10 +382,14 @@ class DesktopRepository( 
        )

        val desk = checkNotNull(desktopData.getDesk(deskId)) { "Did not find desk: $deskId" }

        desk.rightTiledTaskId = taskId

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {

            if (DesktopExperienceFlags.REPOSITORY_BASED_PERSISTENCE.isTrue) {

                updatePersistentRepository(displayId)

            } else {

                updatePersistentRepositoryForDesk(deskId)

            }

        }

    }



    /** Gets a registered left tiled task to desktop state or returns null. */

    fun getLeftTiledTask(deskId: Int): Int? {
@@ -397,20 +408,28 @@ class DesktopRepository( 
        val desk = desktopData.getDesk(deskId)

        if (desk == null) return

        desk.leftTiledTaskId = null

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {

            if (DesktopExperienceFlags.REPOSITORY_BASED_PERSISTENCE.isTrue) {

                updatePersistentRepository(displayId)

            } else {

                updatePersistentRepositoryForDesk(deskId)

            }

        }

    }



    fun removeRightTiledTaskFromDesk(displayId: Int, deskId: Int) {

        logD("removeRightTiledTaskFromDesk for displayId=%d", displayId)

        val desk = desktopData.getDesk(deskId)

        if (desk == null) return

        desk.rightTiledTaskId = null

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {

            if (DesktopExperienceFlags.REPOSITORY_BASED_PERSISTENCE.isTrue) {

                updatePersistentRepository(displayId)

            } else {

                updatePersistentRepositoryForDesk(deskId)

            }

        }

    }



    /** Returns the id of the active desk in the given display, if any. */

    fun getActiveDeskId(displayId: Int): Int? = desktopData.getActiveDesk(displayId)?.deskId
@@ -781,7 +800,7 @@ class DesktopRepository( 
                desk.visibleTasks,

            )

            notifyVisibleTaskListeners(displayId, newCount)

            if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {

            if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {

                updatePersistentRepository(displayId)

            }

        }
@@ -913,7 +932,7 @@ class DesktopRepository( 
        }

        desk.freeformTasksInZOrder.add(0, taskId)

        if (!bounds.isEmpty) desk.boundsByTaskId[taskId] = bounds

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {

            // TODO: can probably just update the desk.

            updatePersistentRepository(displayId)

        }
@@ -947,10 +966,14 @@ class DesktopRepository( 
        desktopData.getDesk(deskId)?.minimizedTasks?.add(taskId)

            ?: logD("Minimize task: No active desk found for task: taskId=%d", taskId)

        updateTaskInDesk(displayId, deskId, taskId, isVisible = false, taskBounds = null)

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {

            if (DesktopExperienceFlags.REPOSITORY_BASED_PERSISTENCE.isTrue) {

                updatePersistentRepository(displayId)

            } else {

                updatePersistentRepositoryForDesk(deskId)

            }

        }

    }



    /**

     * Unminimizes the task for [taskId] and [displayId].
@@ -1024,7 +1047,11 @@ class DesktopRepository( 
        removeActiveTaskFromDesk(deskId = deskId, taskId = taskId)

        removeVisibleTaskFromDesk(deskId = deskId, taskId = taskId)

        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {

            updatePersistentRepositoryForDesk(desk.deskId)

            if (DesktopExperienceFlags.REPOSITORY_BASED_PERSISTENCE.isTrue) {

                updatePersistentRepository(desk.displayId)

            } else {

                updatePersistentRepositoryForDesk(deskId)

            }

        }

    }
@@ -1061,8 +1088,12 @@ class DesktopRepository( 
            DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue &&

                DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue

        ) {

            if (DesktopExperienceFlags.REPOSITORY_BASED_PERSISTENCE.isTrue) {

                updatePersistentRepository(desk.displayId)

            } else {

                removeDeskFromPersistentRepository(desk)

            }

        }

        return activeTasks

    }
@@ -1131,19 +1162,47 @@ class DesktopRepository( 
            }

            .toTypedArray()



    /** TODO: b/389960283 - consider updating only the changing desks. */

    private fun updatePersistentRepository(displayId: Int) {

        val desks = desktopData.desksSequence(displayId).map { desk -> desk.deepCopy() }.toList()

        logD("updatePersistentRepository: displayId=%d", displayId)

        if (displayId == INVALID_DISPLAY) return



        val desks = desktopData.desksSequence(displayId).map { it.deepCopy() }.toList()

        if (desks.isEmpty()) {

            logD("updatePersistentRepository: no desks found for displayId=%d, skipping", displayId)

            return

        }

        if (DesktopExperienceFlags.REPOSITORY_BASED_PERSISTENCE.isTrue) {

            mainCoroutineScope.launch {

                try {

                    logD("updatePersistentRepository user=%d display=%d", userId, displayId)

                    persistentRepository.addOrUpdateRepository(userId, desks)

                } catch (exception: Exception) {

                    logE(

                        "An exception occurred while updating the persistent repository \n%s",

                        exception.stackTrace,

                    )

                }

            }

        } else {

            mainCoroutineScope.launch {

                desks.forEach { desk -> updatePersistentRepositoryForDesk(desk) }

            }

        }

    }



    @Deprecated(

        "Use updatePersistentRepository() instead.",

        ReplaceWith("updatePersistentRepository()"),

    )

    private fun updatePersistentRepositoryForDesk(deskId: Int) {

        val desk = desktopData.getDesk(deskId)?.deepCopy() ?: return

        mainCoroutineScope.launch { updatePersistentRepositoryForDesk(desk) }

    }



    @Deprecated(

        "Use updatePersistentRepository() instead.",

        ReplaceWith("updatePersistentRepository()"),

    )

    private suspend fun updatePersistentRepositoryForDesk(desk: Desk) {

        try {

            persistentRepository.addOrUpdateDesktop(
@@ -1163,6 +1222,10 @@ class DesktopRepository( 
        }

    }



    @Deprecated(

        "Use updatePersistentRepository() instead.",

        ReplaceWith("updatePersistentRepository()"),

    )

    private fun removeDeskFromPersistentRepository(desk: Desk) {

        mainCoroutineScope.launch {

            try {

libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/data/persistence/DesktopPersistentRepository.kt

+55 −2
Original line number Diff line number Diff line
@@ -27,6 +27,7 @@ import androidx.datastore.core.Serializer 
import androidx.datastore.core.handlers.ReplaceFileCorruptionHandler

import androidx.datastore.dataStoreFile

import com.android.framework.protobuf.InvalidProtocolBufferException

import com.android.wm.shell.desktopmode.data.Desk

import com.android.wm.shell.shared.annotations.ShellBackgroundThread

import java.io.IOException

import java.io.InputStream
@@ -107,7 +108,54 @@ class DesktopPersistentRepository(private val dataStore: DataStore<DesktopPersis 
            null

        }



    suspend fun addOrUpdateRepository(userId: Int, desks: List<Desk>) {

        try {

            dataStore.updateData { persistentRepositories: DesktopPersistentRepositories ->

                val currentRepository =

                    persistentRepositories.getDesktopRepoByUserOrDefault(

                        userId,

                        DesktopRepositoryState.getDefaultInstance(),

                    )



                val currentUserRepoBuilder = currentRepository.toBuilder()

                desks.forEach { desk ->

                    if (isEmptyDesk(desk.freeformTasksInZOrder)) {

                        currentUserRepoBuilder.removeDesktop(desk.deskId)

                    } else {

                        val updatedDesktop =

                            getDesktop(currentRepository, desk.deskId, desk.displayId)

                                .toBuilder()

                                .updateTaskStates(

                                    desk.visibleTasks,

                                    desk.minimizedTasks,

                                    desk.freeformTasksInZOrder,

                                    desk.leftTiledTaskId,

                                    desk.rightTiledTaskId,

                                )

                                .updateZOrder(desk.freeformTasksInZOrder)

                                .build()



                        currentUserRepoBuilder.putDesktop(desk.deskId, updatedDesktop)

                    }

                }



                persistentRepositories

                    .toBuilder()

                    .putDesktopRepoByUser(userId, currentUserRepoBuilder.build())

                    .build()

            }

        } catch (exception: Exception) {

            Log.e(

                TAG,

                "Error in updating desktop mode related data, data is " +

                    "stored in a file named $DESKTOP_REPOSITORIES_DATASTORE_FILE",

                exception,

            )

        }

    }



    /** Adds or updates a desktop stored in the datastore */

    @Deprecated("Use addOrUpdateRepository() instead.", ReplaceWith("addOrUpdateRepository()"))

    suspend fun addOrUpdateDesktop(

        userId: Int,

        desktopId: Int = 0,
@@ -169,6 +217,7 @@ class DesktopPersistentRepository(private val dataStore: DataStore<DesktopPersis 
    }



    /** Removes the desktop from the persistent repository. */

    @Deprecated("Use addOrUpdateRepository() instead.", ReplaceWith("addOrUpdateRepository()"))

    suspend fun removeDesktop(userId: Int, desktopId: Int) {

        try {

            dataStore.updateData { persistentRepositories: DesktopPersistentRepositories ->
@@ -212,11 +261,15 @@ class DesktopPersistentRepository(private val dataStore: DataStore<DesktopPersis 
        }

    }



    private fun getDesktop(currentRepository: DesktopRepositoryState, desktopId: Int): Desktop =

    private fun getDesktop(

        currentRepository: DesktopRepositoryState,

        desktopId: Int,

        displayId: Int = DEFAULT_DISPLAY,

    ): Desktop =

        // If there are no desktops set up, create one on the default display

        currentRepository.getDesktopOrDefault(

            desktopId,

            Desktop.newBuilder().setDesktopId(desktopId).setDisplayId(DEFAULT_DISPLAY).build(),

            Desktop.newBuilder().setDesktopId(desktopId).setDisplayId(displayId).build(),

        )



    private fun isEmptyDesk(freeformTasksInZOrder: ArrayList<Int>) = freeformTasksInZOrder.isEmpty()

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/data/DesktopRepositoryTest.kt

+231 −9

File changed.

Preview size limit exceeded, changes collapsed.

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/data/persistence/DesktopPersistentRepositoryTest.kt

+180 −0
Original line number Diff line number Diff line
@@ -30,6 +30,7 @@ import androidx.test.platform.app.InstrumentationRegistry 
import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE

import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE

import com.android.wm.shell.ShellTestCase

import com.android.wm.shell.desktopmode.data.Desk

import com.google.common.truth.Truth.assertThat

import java.io.File

import kotlinx.coroutines.CoroutineScope
@@ -37,7 +38,9 @@ import kotlinx.coroutines.Dispatchers 
import kotlinx.coroutines.ExperimentalCoroutinesApi

import kotlinx.coroutines.SupervisorJob

import kotlinx.coroutines.cancel

import kotlinx.coroutines.launch

import kotlinx.coroutines.test.StandardTestDispatcher

import kotlinx.coroutines.test.advanceUntilIdle

import kotlinx.coroutines.test.runTest

import kotlinx.coroutines.test.setMain

import org.junit.After
@@ -292,6 +295,183 @@ class DesktopPersistentRepositoryTest : ShellTestCase() { 
        }

    }



    @Test

    fun addOrUpdateRepository_addNewTaskToDesktop() {

        runTest(StandardTestDispatcher()) {

            // Create a basic repository state

            val task = createDesktopTask(1)

            val desktopPersistentRepositories = createRepositoryWithOneDesk(ArrayList(listOf(task)))

            testDatastore.updateData { desktopPersistentRepositories }

            // Create a new state to be initialized

            val visibleTasks = ArraySet(listOf(1, 2))

            val minimizedTasks = ArraySet<Int>()

            val freeformTasksInZOrder = ArrayList(listOf(2, 1))

            val desk =

                Desk(

                    deskId = DEFAULT_DESKTOP_ID,

                    displayId = DEFAULT_DISPLAY,

                    visibleTasks = visibleTasks,

                    minimizedTasks = minimizedTasks,

                    freeformTasksInZOrder = freeformTasksInZOrder,

                )



            // Update with new state

            datastoreRepository.addOrUpdateRepository(

                userId = DEFAULT_USER_ID,

                desks = listOf(desk),

            )



            val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)

            assertThat(actualDesktop?.tasksByTaskIdMap).hasSize(2)

            assertThat(actualDesktop?.getZOrderedTasks(0)).isEqualTo(2)

        }

    }



    @Test

    fun addOrUpdateRepository_tiledTasks_addsTiledTasksToDesktop() {

        runTest(StandardTestDispatcher()) {

            // Create a basic repository state

            val task = createDesktopTask(1)

            val desktopPersistentRepositories = createRepositoryWithOneDesk(ArrayList(listOf(task)))

            testDatastore.updateData { desktopPersistentRepositories }

            // Create a new state to be initialized

            val visibleTasks = ArraySet(listOf(1, 2))

            val freeformTasksInZOrder = ArrayList(listOf(1, 2))

            val desk =

                Desk(

                    deskId = DEFAULT_DESKTOP_ID,

                    displayId = DEFAULT_DISPLAY,

                    visibleTasks = visibleTasks,

                    freeformTasksInZOrder = freeformTasksInZOrder,

                    leftTiledTaskId = 1,

                )



            // Update with new state

            datastoreRepository.addOrUpdateRepository(

                userId = DEFAULT_USER_ID,

                desks = listOf(desk),

            )



            var actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)

            assertThat(actualDesktop?.tasksByTaskIdMap?.get(1)?.desktopTaskTilingState)

                .isEqualTo(DesktopTaskTilingState.LEFT)

            assertThat(actualDesktop?.tasksByTaskIdMap?.get(2)?.desktopTaskTilingState)

                .isEqualTo(DesktopTaskTilingState.NONE)



            desk.rightTiledTaskId = 2

            desk.leftTiledTaskId = null



            // Update with new state

            datastoreRepository.addOrUpdateRepository(

                userId = DEFAULT_USER_ID,

                desks = listOf(desk),

            )



            actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)

            assertThat(actualDesktop?.tasksByTaskIdMap?.get(1)?.desktopTaskTilingState)

                .isEqualTo(DesktopTaskTilingState.NONE)

            assertThat(actualDesktop?.tasksByTaskIdMap?.get(2)?.desktopTaskTilingState)

                .isEqualTo(DesktopTaskTilingState.RIGHT)

        }

    }



    @Test

    fun addOrUpdateRepository_changeTaskStateToMinimize_taskStateIsMinimized() {

        runTest(StandardTestDispatcher()) {

            val task = createDesktopTask(1)

            val desktopPersistentRepositories = createRepositoryWithOneDesk(ArrayList(listOf(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))

            val desk =

                Desk(

                    deskId = DEFAULT_DESKTOP_ID,

                    displayId = DEFAULT_DISPLAY,

                    visibleTasks = visibleTasks,

                    minimizedTasks = minimizedTasks,

                    freeformTasksInZOrder = freeformTasksInZOrder,

                )



            // Update with new state

            datastoreRepository.addOrUpdateRepository(

                userId = DEFAULT_USER_ID,

                desks = listOf(desk),

            )



            val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)

            assertThat(actualDesktop?.tasksByTaskIdMap?.get(task.taskId)?.desktopTaskState)

                .isEqualTo(DesktopTaskState.MINIMIZED)

        }

    }



    @Test

    fun addOrUpdateRepository_removeMultipleDesks_removesAllDesks() {

        runTest(StandardTestDispatcher()) {

            val task = createDesktopTask(1)

            val desktopPersistentRepositories = createRepositoryWithOneDesk(ArrayList(listOf(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))

            val desk1 =

                Desk(

                    deskId = DEFAULT_DESKTOP_ID,

                    displayId = DEFAULT_DISPLAY,

                    visibleTasks = visibleTasks,

                    minimizedTasks = minimizedTasks,

                    freeformTasksInZOrder = freeformTasksInZOrder,

                )

            val desk2 =

                Desk(

                    deskId = DEFAULT_DESKTOP_ID + 1,

                    displayId = DEFAULT_DISPLAY,

                    visibleTasks = visibleTasks,

                    minimizedTasks = minimizedTasks,

                    freeformTasksInZOrder = freeformTasksInZOrder,

                )

            val desk3 =

                Desk(

                    deskId = DEFAULT_DESKTOP_ID + 3,

                    displayId = DEFAULT_DISPLAY,

                    visibleTasks = visibleTasks,

                    minimizedTasks = minimizedTasks,

                    freeformTasksInZOrder = freeformTasksInZOrder,

                )

            datastoreRepository.addOrUpdateRepository(

                userId = DEFAULT_USER_ID,

                desks = listOf(desk1, desk2, desk3),

            )



            // Back to back removals

            launch {

                datastoreRepository.addOrUpdateRepository(

                    userId = DEFAULT_USER_ID,

                    desks = listOf(desk1, desk3),

                )

            }

            launch {

                datastoreRepository.addOrUpdateRepository(

                    userId = DEFAULT_USER_ID,

                    desks = listOf(desk3),

                )

            }

            launch {

                datastoreRepository.addOrUpdateRepository(

                    userId = DEFAULT_USER_ID,

                    desks = listOf(),

                )

            }

            advanceUntilIdle()



            val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)

            assertThat(actualDesktop?.tasksByTaskIdMap?.get(task.taskId)?.desktopTaskState)

                .isEqualTo(DesktopTaskState.MINIMIZED)

        }

    }



    private companion object {

        const val DESKTOP_REPOSITORY_STATES_DATASTORE_TEST_FILE = "desktop_repo_test.pb"

        const val DEFAULT_USER_ID = 1000