Improve performance of thumbnail retrieval - hot taskData flow

Bug: 357542209
Test: TasksRepositoryTest, manual with reduced speed video
Flag: com.android.launcher3.enable_refactor_task_thumbnail
Change-Id: Ie94daf61d8e41c4b9c12d67f5f45166104879325
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 6acc940..6f9d157 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -17,65 +17,104 @@
 package com.android.quickstep.recents.data
 
 import android.graphics.drawable.Drawable
+import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.quickstep.task.thumbnail.data.TaskIconDataSource
 import com.android.quickstep.task.thumbnail.data.TaskThumbnailDataSource
 import com.android.quickstep.util.GroupTask
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import kotlin.coroutines.resume
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.distinctUntilChangedBy
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlinx.coroutines.withContext
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class TasksRepository(
     private val recentsModel: RecentTasksDataSource,
     private val taskThumbnailDataSource: TaskThumbnailDataSource,
     private val taskIconDataSource: TaskIconDataSource,
+    recentsCoroutineScope: CoroutineScope,
+    private val dispatcherProvider: DispatcherProvider,
 ) : RecentTasksRepository {
     private val groupedTaskData = MutableStateFlow(emptyList<GroupTask>())
-    private val _taskData =
-        groupedTaskData.map { groupTaskList -> groupTaskList.flatMap { it.tasks } }
     private val visibleTaskIds = MutableStateFlow(emptySet<Int>())
     private val thumbnailOverride = MutableStateFlow(mapOf<Int, ThumbnailData>())
 
-    private val taskData: Flow<List<Task>> =
-        combine(_taskData, getThumbnailQueryResults(), getIconQueryResults(), thumbnailOverride) {
-            tasks,
-            thumbnailQueryResults,
-            iconQueryResults,
-            thumbnailOverride ->
-            tasks.forEach { task ->
-                // Add retrieved thumbnails + remove unnecessary thumbnails (e.g. invisible)
-                task.thumbnail =
-                    thumbnailOverride[task.key.id] ?: thumbnailQueryResults[task.key.id]
-
-                // TODO(b/352331675) don't load icons for DesktopTaskView
-                // Add retrieved icons + remove unnecessary icons
-                task.icon = iconQueryResults[task.key.id]?.icon
-                task.titleDescription = iconQueryResults[task.key.id]?.contentDescription
-                task.title = iconQueryResults[task.key.id]?.title
-            }
-            tasks
+    private val taskData =
+        groupedTaskData.map { groupTaskList -> groupTaskList.flatMap { it.tasks } }
+    private val visibleTasks =
+        combine(taskData, visibleTaskIds) { tasks, visibleIds ->
+            tasks.filter { it.key.id in visibleIds }
         }
 
+    private val iconQueryResults: Flow<Map<Int, TaskIconQueryResponse?>> =
+        visibleTasks
+            .map { visibleTasksList -> visibleTasksList.map(::getIconDataRequest) }
+            .flatMapLatest { iconRequestFlows: List<IconDataRequest> ->
+                if (iconRequestFlows.isEmpty()) {
+                    flowOf(emptyMap())
+                } else {
+                    combine(iconRequestFlows) { it.toMap() }
+                }
+            }
+            .distinctUntilChanged()
+
+    private val thumbnailQueryResults: Flow<Map<Int, ThumbnailData?>> =
+        visibleTasks
+            .map { visibleTasksList -> visibleTasksList.map(::getThumbnailDataRequest) }
+            .flatMapLatest { thumbnailRequestFlows: List<ThumbnailDataRequest> ->
+                if (thumbnailRequestFlows.isEmpty()) {
+                    flowOf(emptyMap())
+                } else {
+                    combine(thumbnailRequestFlows) { it.toMap() }
+                }
+            }
+            .distinctUntilChanged()
+
+    private val augmentedTaskData: Flow<List<Task>> =
+        combine(taskData, thumbnailQueryResults, iconQueryResults, thumbnailOverride) {
+                tasks,
+                thumbnailQueryResults,
+                iconQueryResults,
+                thumbnailOverride ->
+                tasks.onEach { task ->
+                    // Add retrieved thumbnails + remove unnecessary thumbnails (e.g. invisible)
+                    task.thumbnail =
+                        thumbnailOverride[task.key.id] ?: thumbnailQueryResults[task.key.id]
+
+                    // TODO(b/352331675) don't load icons for DesktopTaskView
+                    // Add retrieved icons + remove unnecessary icons
+                    val iconQueryResult = iconQueryResults[task.key.id]
+                    task.icon = iconQueryResult?.icon
+                    task.titleDescription = iconQueryResult?.contentDescription
+                    task.title = iconQueryResult?.title
+                }
+            }
+            .flowOn(dispatcherProvider.io)
+            .shareIn(recentsCoroutineScope, SharingStarted.WhileSubscribed(), replay = 1)
+
     override fun getAllTaskData(forceRefresh: Boolean): Flow<List<Task>> {
         if (forceRefresh) {
             recentsModel.getTasks { groupedTaskData.value = it }
         }
-        return taskData
+        return augmentedTaskData
     }
 
     override fun getTaskDataById(taskId: Int): Flow<Task?> =
-        taskData.map { taskList -> taskList.firstOrNull { it.key.id == taskId } }
+        augmentedTaskData.map { taskList -> taskList.firstOrNull { it.key.id == taskId } }
 
     override fun getThumbnailById(taskId: Int): Flow<ThumbnailData?> =
         getTaskDataById(taskId).map { it?.thumbnail }.distinctUntilChangedBy { it?.snapshotId }
@@ -94,41 +133,19 @@
     }
 
     /** Flow wrapper for [TaskThumbnailDataSource.getThumbnailInBackground] api */
-    private fun getThumbnailDataRequest(task: Task): ThumbnailDataRequest =
-        flow {
-                emit(task.key.id to task.thumbnail)
-                val thumbnailDataResult: ThumbnailData? =
-                    suspendCancellableCoroutine { continuation ->
-                        val cancellableTask =
-                            taskThumbnailDataSource.getThumbnailInBackground(task) {
-                                continuation.resume(it)
-                            }
-                        continuation.invokeOnCancellation { cancellableTask?.cancel() }
-                    }
-                emit(task.key.id to thumbnailDataResult)
+    private fun getThumbnailDataRequest(task: Task): ThumbnailDataRequest = flow {
+        emit(task.key.id to task.thumbnail)
+        val thumbnailDataResult: ThumbnailData? =
+            withContext(dispatcherProvider.main) {
+                suspendCancellableCoroutine { continuation ->
+                    val cancellableTask =
+                        taskThumbnailDataSource.getThumbnailInBackground(task) {
+                            continuation.resume(it)
+                        }
+                    continuation.invokeOnCancellation { cancellableTask?.cancel() }
+                }
             }
-            .distinctUntilChanged()
-
-    /**
-     * This is a Flow that makes a query for thumbnail data to the [taskThumbnailDataSource] for
-     * each visible task. It then collects the responses and returns them in a Map as soon as they
-     * are available.
-     */
-    private fun getThumbnailQueryResults(): Flow<Map<Int, ThumbnailData?>> {
-        val visibleTasks =
-            combine(_taskData, visibleTaskIds) { tasks, visibleIds ->
-                tasks.filter { it.key.id in visibleIds }
-            }
-        val visibleThumbnailDataRequests: Flow<List<ThumbnailDataRequest>> =
-            visibleTasks.map { visibleTasksList -> visibleTasksList.map(::getThumbnailDataRequest) }
-        return visibleThumbnailDataRequests.flatMapLatest {
-            thumbnailRequestFlows: List<ThumbnailDataRequest> ->
-            if (thumbnailRequestFlows.isEmpty()) {
-                flowOf(emptyMap())
-            } else {
-                combine(thumbnailRequestFlows) { it.toMap() }
-            }
-        }
+        emit(task.key.id to thumbnailDataResult)
     }
 
     /** Flow wrapper for [TaskThumbnailDataSource.getThumbnailInBackground] api */
@@ -136,43 +153,29 @@
         flow {
                 emit(task.key.id to task.getTaskIconQueryResponse())
                 val iconDataResponse: TaskIconQueryResponse? =
-                    suspendCancellableCoroutine { continuation ->
-                        val cancellableTask =
-                            taskIconDataSource.getIconInBackground(task) {
-                                icon,
-                                contentDescription,
-                                title ->
-                                icon.constantState?.let {
-                                    continuation.resume(
-                                        TaskIconQueryResponse(
-                                            it.newDrawable().mutate(),
-                                            contentDescription,
-                                            title
+                    withContext(dispatcherProvider.main) {
+                        suspendCancellableCoroutine { continuation ->
+                            val cancellableTask =
+                                taskIconDataSource.getIconInBackground(task) {
+                                    icon,
+                                    contentDescription,
+                                    title ->
+                                    icon.constantState?.let {
+                                        continuation.resume(
+                                            TaskIconQueryResponse(
+                                                it.newDrawable().mutate(),
+                                                contentDescription,
+                                                title
+                                            )
                                         )
-                                    )
+                                    }
                                 }
-                            }
-                        continuation.invokeOnCancellation { cancellableTask?.cancel() }
+                            continuation.invokeOnCancellation { cancellableTask?.cancel() }
+                        }
                     }
                 emit(task.key.id to iconDataResponse)
             }
             .distinctUntilChanged()
-
-    private fun getIconQueryResults(): Flow<Map<Int, TaskIconQueryResponse?>> {
-        val visibleTasks =
-            combine(_taskData, visibleTaskIds) { tasks, visibleIds ->
-                tasks.filter { it.key.id in visibleIds }
-            }
-        val visibleIconDataRequests: Flow<List<IconDataRequest>> =
-            visibleTasks.map { visibleTasksList -> visibleTasksList.map(::getIconDataRequest) }
-        return visibleIconDataRequests.flatMapLatest { iconRequestFlows: List<IconDataRequest> ->
-            if (iconRequestFlows.isEmpty()) {
-                flowOf(emptyMap())
-            } else {
-                combine(iconRequestFlows) { it.toMap() }
-            }
-        }
-    }
 }
 
 data class TaskIconQueryResponse(
diff --git a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
index eba7688..d8156b1 100644
--- a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
+++ b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
@@ -19,6 +19,7 @@
 import android.content.Context
 import android.util.Log
 import android.view.View
+import com.android.launcher3.util.coroutines.ProductionDispatchers
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.recents.data.RecentTasksRepository
 import com.android.quickstep.recents.data.TasksRepository
@@ -36,6 +37,10 @@
 import com.android.quickstep.task.viewmodel.TaskViewModel
 import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.Task
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
 
 internal typealias RecentsScopeId = String
 
@@ -53,11 +58,20 @@
     private fun startDefaultScope(appContext: Context) {
         createScope(DEFAULT_SCOPE_ID).apply {
             set(RecentsViewData::class.java.simpleName, RecentsViewData())
+            val recentsCoroutineScope =
+                CoroutineScope(SupervisorJob() + Dispatchers.Main + CoroutineName("RecentsView"))
+            set(CoroutineScope::class.java.simpleName, recentsCoroutineScope)
 
             // Create RecentsTaskRepository singleton
             val recentTasksRepository: RecentTasksRepository =
                 with(RecentsModel.INSTANCE.get(appContext)) {
-                    TasksRepository(this, thumbnailCache, iconCache)
+                    TasksRepository(
+                        this,
+                        thumbnailCache,
+                        iconCache,
+                        recentsCoroutineScope,
+                        ProductionDispatchers
+                    )
                 }
             set(RecentTasksRepository::class.java.simpleName, recentTasksRepository)
         }
@@ -137,7 +151,13 @@
             when (modelClass) {
                 RecentTasksRepository::class.java -> {
                     with(RecentsModel.INSTANCE.get(appContext)) {
-                        TasksRepository(this, thumbnailCache, iconCache)
+                        TasksRepository(
+                            this,
+                            thumbnailCache,
+                            iconCache,
+                            get(),
+                            ProductionDispatchers
+                        )
                     }
                 }
                 RecentsViewData::class.java -> RecentsViewData()
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
index b34e156..e6534eb 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -20,6 +20,7 @@
 import android.content.Intent
 import android.graphics.Bitmap
 import android.view.Surface
+import com.android.launcher3.util.TestDispatcherProvider
 import com.android.quickstep.task.thumbnail.TaskThumbnailViewModelTest
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.GroupTask
@@ -27,10 +28,8 @@
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.drop
 import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.flow.toList
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -50,223 +49,232 @@
     private val taskThumbnailDataSource = FakeTaskThumbnailDataSource()
     private val taskIconDataSource = FakeTaskIconDataSource()
 
+    private val dispatcher = UnconfinedTestDispatcher()
+    private val testScope = TestScope(dispatcher)
     private val systemUnderTest =
-        TasksRepository(recentsModel, taskThumbnailDataSource, taskIconDataSource)
+        TasksRepository(
+            recentsModel,
+            taskThumbnailDataSource,
+            taskIconDataSource,
+            testScope.backgroundScope,
+            TestDispatcherProvider(dispatcher)
+        )
 
     @Test
-    fun getAllTaskDataReturnsFlattenedListOfTasks() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
+    fun getAllTaskDataReturnsFlattenedListOfTasks() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
 
-        assertThat(systemUnderTest.getAllTaskData(forceRefresh = true).first()).isEqualTo(tasks)
-    }
-
-    @Test
-    fun getTaskDataByIdReturnsSpecificTask() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-
-        assertThat(systemUnderTest.getTaskDataById(2).first()).isEqualTo(tasks[2])
-    }
-
-    @Test
-    fun setVisibleTasksPopulatesThumbnails() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
-        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        assertThat(systemUnderTest.getTaskDataById(1).drop(1).first()!!.thumbnail!!.thumbnail)
-            .isEqualTo(bitmap1)
-        assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
-            .isEqualTo(bitmap2)
-    }
-
-    @Test
-    fun setVisibleTasksPopulatesIcons() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        systemUnderTest
-            .getTaskDataById(1)
-            .drop(1)
-            .first()!!
-            .assertHasIconDataFromSource(taskIconDataSource)
-        systemUnderTest.getTaskDataById(2).first()!!.assertHasIconDataFromSource(taskIconDataSource)
-    }
-
-    @Test
-    fun changingVisibleTasksContainsAlreadyPopulatedThumbnails() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        assertThat(systemUnderTest.getTaskDataById(2).drop(1).first()!!.thumbnail!!.thumbnail)
-            .isEqualTo(bitmap2)
-
-        // Prevent new loading of Bitmaps
-        taskThumbnailDataSource.shouldLoadSynchronously = false
-        systemUnderTest.setVisibleTasks(listOf(2, 3))
-
-        assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
-            .isEqualTo(bitmap2)
-    }
-
-    @Test
-    fun changingVisibleTasksContainsAlreadyPopulatedIcons() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-
-        // .drop(1) to ignore initial null content before from icon was loaded.
-        systemUnderTest
-            .getTaskDataById(2)
-            .drop(1)
-            .first()!!
-            .assertHasIconDataFromSource(taskIconDataSource)
-
-        // Prevent new loading of Drawables
-        taskThumbnailDataSource.shouldLoadSynchronously = false
-        systemUnderTest.setVisibleTasks(listOf(2, 3))
-
-        systemUnderTest.getTaskDataById(2).first()!!.assertHasIconDataFromSource(taskIconDataSource)
-    }
-
-    @Test
-    fun retrievedImagesAreDiscardedWhenTaskBecomesInvisible() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        val task2 = systemUnderTest.getTaskDataById(2).drop(1).first()!!
-        assertThat(task2.thumbnail!!.thumbnail).isEqualTo(bitmap2)
-        task2.assertHasIconDataFromSource(taskIconDataSource)
-
-        // Prevent new loading of Bitmaps
-        taskThumbnailDataSource.shouldLoadSynchronously = false
-        taskIconDataSource.shouldLoadSynchronously = false
-        systemUnderTest.setVisibleTasks(listOf(0, 1))
-
-        val task2AfterVisibleTasksChanged = systemUnderTest.getTaskDataById(2).first()!!
-        assertThat(task2AfterVisibleTasksChanged.thumbnail).isNull()
-        assertThat(task2AfterVisibleTasksChanged.icon).isNull()
-        assertThat(task2AfterVisibleTasksChanged.titleDescription).isNull()
-        assertThat(task2AfterVisibleTasksChanged.title).isNull()
-    }
-
-    @Test
-    fun retrievedThumbnailsCauseEmissionOnTaskDataFlow() = runTest {
-        // Setup fakes
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-        taskThumbnailDataSource.shouldLoadSynchronously = false
-
-        // Setup TasksRepository
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-
-        // Assert there is no bitmap in first emission
-        val taskFlow = systemUnderTest.getTaskDataById(2)
-        val taskFlowValuesList = mutableListOf<Task?>()
-        backgroundScope.launch(UnconfinedTestDispatcher(testScheduler)) {
-            taskFlow.toList(taskFlowValuesList)
+            assertThat(systemUnderTest.getAllTaskData(forceRefresh = true).first()).isEqualTo(tasks)
         }
-        assertThat(taskFlowValuesList[0]!!.thumbnail).isNull()
-
-        // Simulate bitmap loading after first emission
-        taskThumbnailDataSource.taskIdToUpdatingTask.getValue(2).invoke()
-
-        // Check for second emission
-        assertThat(taskFlowValuesList[1]!!.thumbnail!!.thumbnail).isEqualTo(bitmap2)
-    }
 
     @Test
-    fun addThumbnailOverrideOverrideThumbnails() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
-        val thumbnailOverride2 = createThumbnailData()
-        systemUnderTest.getAllTaskData(forceRefresh = true)
+    fun getTaskDataByIdReturnsSpecificTask() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            systemUnderTest.getAllTaskData(forceRefresh = true)
 
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
-
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
-            .isEqualTo(bitmap1)
-        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail)
-            .isEqualTo(thumbnailOverride2.thumbnail)
-    }
+            assertThat(systemUnderTest.getTaskDataById(2).first()).isEqualTo(tasks[2])
+        }
 
     @Test
-    fun addThumbnailOverrideMultipleOverrides() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val thumbnailOverride1 = createThumbnailData()
-        val thumbnailOverride2 = createThumbnailData()
-        val thumbnailOverride3 = createThumbnailData()
-        systemUnderTest.getAllTaskData(forceRefresh = true)
+    fun setVisibleTasksPopulatesThumbnails() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+            systemUnderTest.getAllTaskData(forceRefresh = true)
 
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(1 to thumbnailOverride1))
-        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
-        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride3))
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
 
-        assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail)
-            .isEqualTo(thumbnailOverride1.thumbnail)
-        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail)
-            .isEqualTo(thumbnailOverride3.thumbnail)
-    }
+            assertThat(systemUnderTest.getTaskDataById(1).first()!!.thumbnail!!.thumbnail)
+                .isEqualTo(bitmap1)
+            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
+                .isEqualTo(bitmap2)
+        }
 
     @Test
-    fun addThumbnailOverrideClearedWhenTaskBecomeInvisible() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-        val thumbnailOverride1 = createThumbnailData()
-        val thumbnailOverride2 = createThumbnailData()
-        systemUnderTest.getAllTaskData(forceRefresh = true)
+    fun setVisibleTasksPopulatesIcons() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            systemUnderTest.getAllTaskData(forceRefresh = true)
 
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(1 to thumbnailOverride1))
-        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
-        // Making task 2 invisible and visible again should clear the override
-        systemUnderTest.setVisibleTasks(listOf(1))
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
 
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail)
-            .isEqualTo(thumbnailOverride1.thumbnail)
-        assertThat(systemUnderTest.getThumbnailById(2).drop(1).first()!!.thumbnail)
-            .isEqualTo(bitmap2)
-    }
+            systemUnderTest
+                .getTaskDataById(1)
+                .first()!!
+                .assertHasIconDataFromSource(taskIconDataSource)
+            systemUnderTest
+                .getTaskDataById(2)
+                .first()!!
+                .assertHasIconDataFromSource(taskIconDataSource)
+        }
 
     @Test
-    fun addThumbnailOverrideDoesNotOverrideInvisibleTasks() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
-        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-        val thumbnailOverride = createThumbnailData()
-        systemUnderTest.getAllTaskData(forceRefresh = true)
+    fun changingVisibleTasksContainsAlreadyPopulatedThumbnails() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+            systemUnderTest.getAllTaskData(forceRefresh = true)
 
-        systemUnderTest.setVisibleTasks(listOf(1))
-        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride))
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
 
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
-            .isEqualTo(bitmap1)
-        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail).isEqualTo(bitmap2)
-    }
+            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
+                .isEqualTo(bitmap2)
+
+            // Prevent new loading of Bitmaps
+            taskThumbnailDataSource.shouldLoadSynchronously = false
+            systemUnderTest.setVisibleTasks(listOf(2, 3))
+
+            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
+                .isEqualTo(bitmap2)
+        }
+
+    @Test
+    fun changingVisibleTasksContainsAlreadyPopulatedIcons() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+            systemUnderTest
+                .getTaskDataById(2)
+                .first()!!
+                .assertHasIconDataFromSource(taskIconDataSource)
+
+            // Prevent new loading of Drawables
+            taskThumbnailDataSource.shouldLoadSynchronously = false
+            systemUnderTest.setVisibleTasks(listOf(2, 3))
+
+            systemUnderTest
+                .getTaskDataById(2)
+                .first()!!
+                .assertHasIconDataFromSource(taskIconDataSource)
+        }
+
+    @Test
+    fun retrievedImagesAreDiscardedWhenTaskBecomesInvisible() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+            val task2 = systemUnderTest.getTaskDataById(2).first()!!
+            assertThat(task2.thumbnail!!.thumbnail).isEqualTo(bitmap2)
+            task2.assertHasIconDataFromSource(taskIconDataSource)
+
+            // Prevent new loading of Bitmaps
+            taskThumbnailDataSource.shouldLoadSynchronously = false
+            taskIconDataSource.shouldLoadSynchronously = false
+            systemUnderTest.setVisibleTasks(listOf(0, 1))
+
+            val task2AfterVisibleTasksChanged = systemUnderTest.getTaskDataById(2).first()!!
+            assertThat(task2AfterVisibleTasksChanged.thumbnail).isNull()
+            assertThat(task2AfterVisibleTasksChanged.icon).isNull()
+            assertThat(task2AfterVisibleTasksChanged.titleDescription).isNull()
+            assertThat(task2AfterVisibleTasksChanged.title).isNull()
+        }
+
+    @Test
+    fun retrievedThumbnailsCauseEmissionOnTaskDataFlow() =
+        testScope.runTest {
+            // Setup fakes
+            recentsModel.seedTasks(defaultTaskList)
+            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+            taskThumbnailDataSource.shouldLoadSynchronously = false
+
+            // Setup TasksRepository
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+            // Assert there is no bitmap in first emission
+            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail).isNull()
+
+            // Simulate bitmap loading after first emission
+            taskThumbnailDataSource.taskIdToUpdatingTask.getValue(2).invoke()
+
+            // Check for second emission
+            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
+                .isEqualTo(bitmap2)
+        }
+
+    @Test
+    fun addThumbnailOverrideOverrideThumbnails() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+            val thumbnailOverride2 = createThumbnailData()
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
+            systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
+
+            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail).isEqualTo(bitmap1)
+            assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail)
+                .isEqualTo(thumbnailOverride2.thumbnail)
+        }
+
+    @Test
+    fun addThumbnailOverrideMultipleOverrides() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            val thumbnailOverride1 = createThumbnailData()
+            val thumbnailOverride2 = createThumbnailData()
+            val thumbnailOverride3 = createThumbnailData()
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
+            systemUnderTest.addOrUpdateThumbnailOverride(mapOf(1 to thumbnailOverride1))
+            systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
+            systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride3))
+
+            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail)
+                .isEqualTo(thumbnailOverride1.thumbnail)
+            assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail)
+                .isEqualTo(thumbnailOverride3.thumbnail)
+        }
+
+    @Test
+    fun addThumbnailOverrideClearedWhenTaskBecomeInvisible() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+            val thumbnailOverride1 = createThumbnailData()
+            val thumbnailOverride2 = createThumbnailData()
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
+            systemUnderTest.addOrUpdateThumbnailOverride(mapOf(1 to thumbnailOverride1))
+            systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
+            // Making task 2 invisible and visible again should clear the override
+            systemUnderTest.setVisibleTasks(listOf(1))
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail)
+                .isEqualTo(thumbnailOverride1.thumbnail)
+            assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail).isEqualTo(bitmap2)
+        }
+
+    @Test
+    fun addThumbnailOverrideDoesNotOverrideInvisibleTasks() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+            val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+            val thumbnailOverride = createThumbnailData()
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+
+            systemUnderTest.setVisibleTasks(listOf(1))
+            systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride))
+            systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail).isEqualTo(bitmap1)
+            assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail).isEqualTo(bitmap2)
+        }
 
     private fun createTaskWithId(taskId: Int) =
         Task(Task.TaskKey(taskId, 0, Intent(), ComponentName("", ""), 0, 2000))
diff --git a/src/com/android/launcher3/util/coroutines/DispatcherProvider.kt b/src/com/android/launcher3/util/coroutines/DispatcherProvider.kt
new file mode 100644
index 0000000..e9691a8
--- /dev/null
+++ b/src/com/android/launcher3/util/coroutines/DispatcherProvider.kt
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2024 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.launcher3.util.coroutines
+
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
+
+interface DispatcherProvider {
+    val default: CoroutineDispatcher
+    val io: CoroutineDispatcher
+    val main: CoroutineDispatcher
+    val unconfined: CoroutineDispatcher
+}
+
+object ProductionDispatchers : DispatcherProvider {
+    override val default: CoroutineDispatcher = Dispatchers.Default
+    override val io: CoroutineDispatcher = Dispatchers.IO
+    override val main: CoroutineDispatcher = Dispatchers.Main
+    override val unconfined: CoroutineDispatcher = Dispatchers.Unconfined
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/TestDispatcherProvider.kt b/tests/multivalentTests/src/com/android/launcher3/util/TestDispatcherProvider.kt
new file mode 100644
index 0000000..39e1ec5
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/TestDispatcherProvider.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2024 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.launcher3.util
+
+import com.android.launcher3.util.coroutines.DispatcherProvider
+import kotlinx.coroutines.CoroutineDispatcher
+
+class TestDispatcherProvider(testDispatcher: CoroutineDispatcher) : DispatcherProvider {
+    override val default: CoroutineDispatcher = testDispatcher
+    override val io: CoroutineDispatcher = testDispatcher
+    override val main: CoroutineDispatcher = testDispatcher
+    override val unconfined: CoroutineDispatcher = testDispatcher
+}