Merge "[Mobile] Create a child scope for mobile view models" into main

2 files changed, 58 insertions, 21 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModel.kt
index be843ba8699f..deae576662e3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModel.kt
@@ -22,7 +22,6 @@ import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.flags.FeatureFlagsClassic
 import com.android.systemui.statusbar.phone.StatusBarLocation
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
-import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconInteractor
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconsInteractor
 import com.android.systemui.statusbar.pipeline.mobile.ui.MobileViewLogger
 import com.android.systemui.statusbar.pipeline.mobile.ui.VerboseMobileViewLogger
@@ -31,6 +30,8 @@ import com.android.systemui.statusbar.pipeline.shared.ConnectivityConstants
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.flatMapLatest
@@ -58,9 +59,8 @@ constructor(
     private val flags: FeatureFlagsClassic,
     @Application private val scope: CoroutineScope,
 ) {
-    @VisibleForTesting val mobileIconSubIdCache = mutableMapOf<Int, MobileIconViewModel>()
     @VisibleForTesting
-    val mobileIconInteractorSubIdCache = mutableMapOf<Int, MobileIconInteractor>()
+    val reuseCache = mutableMapOf<Int, Pair<MobileIconViewModel, CoroutineScope>>()
 
     val subscriptionIdsFlow: StateFlow<List<Int>> =
         interactor.filteredSubscriptions
@@ -109,24 +109,37 @@ constructor(
     }
 
     private fun commonViewModelForSub(subId: Int): MobileIconViewModelCommon {
-        return mobileIconSubIdCache[subId]
-            ?: MobileIconViewModel(
-                    subId,
-                    interactor.getMobileConnectionInteractorForSubId(subId),
-                    airplaneModeInteractor,
-                    constants,
-                    flags,
-                    scope,
-                )
-                .also { mobileIconSubIdCache[subId] = it }
+        return reuseCache.getOrPut(subId) { createViewModel(subId) }.first
     }
 
-    private fun invalidateCaches(subIds: List<Int>) {
-        val subIdsToRemove = mobileIconSubIdCache.keys.filter { !subIds.contains(it) }
-        subIdsToRemove.forEach { mobileIconSubIdCache.remove(it) }
+    private fun createViewModel(subId: Int): Pair<MobileIconViewModel, CoroutineScope> {
+        // Create a child scope so we can cancel it
+        val vmScope = scope.createChildScope()
+        val vm =
+            MobileIconViewModel(
+                subId,
+                interactor.getMobileConnectionInteractorForSubId(subId),
+                airplaneModeInteractor,
+                constants,
+                flags,
+                vmScope,
+            )
+
+        return Pair(vm, vmScope)
+    }
 
-        mobileIconInteractorSubIdCache.keys
+    private fun CoroutineScope.createChildScope() =
+        CoroutineScope(coroutineContext + Job(coroutineContext[Job]))
+
+    private fun invalidateCaches(subIds: List<Int>) {
+        reuseCache.keys
             .filter { !subIds.contains(it) }
-            .forEach { subId -> mobileIconInteractorSubIdCache.remove(subId) }
+            .forEach { id ->
+                reuseCache
+                    .remove(id)
+                    // Cancel the view model's scope after removing it
+                    ?.second
+                    ?.cancel()
+            }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
index 1f27a289e52c..47899a66b772 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
@@ -38,9 +38,12 @@ import com.android.systemui.statusbar.pipeline.shared.ConnectivityConstants
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
+import junit.framework.Assert.assertFalse
+import junit.framework.Assert.assertTrue
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.isActive
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
@@ -156,14 +159,35 @@ class MobileIconsViewModelTest : SysuiTestCase() {
             val model2 = underTest.viewModelForSub(2, StatusBarLocation.QS)
 
             // Both impls are cached
-            assertThat(underTest.mobileIconSubIdCache)
-                .containsExactly(1, model1.commonImpl, 2, model2.commonImpl)
+            assertThat(underTest.reuseCache.keys).containsExactly(1, 2)
 
             // SUB_1 is removed from the list...
             interactor.filteredSubscriptions.value = listOf(SUB_2)
 
             // ... and dropped from the cache
-            assertThat(underTest.mobileIconSubIdCache).containsExactly(2, model2.commonImpl)
+            assertThat(underTest.reuseCache.keys).containsExactly(2)
+        }
+
+    @Test
+    fun caching_invalidatedViewModelsAreCanceled() =
+        testScope.runTest {
+            // Retrieve models to trigger caching
+            val model1 = underTest.viewModelForSub(1, StatusBarLocation.HOME)
+            val model2 = underTest.viewModelForSub(2, StatusBarLocation.QS)
+
+            var scope1 = underTest.reuseCache[1]?.second
+            var scope2 = underTest.reuseCache[2]?.second
+
+            // Scopes are not canceled
+            assertTrue(scope1!!.isActive)
+            assertTrue(scope2!!.isActive)
+
+            // SUB_1 is removed from the list...
+            interactor.filteredSubscriptions.value = listOf(SUB_2)
+
+            // scope1 is canceled
+            assertFalse(scope1!!.isActive)
+            assertTrue(scope2!!.isActive)
         }
 
     @Test