Improve device repo list caching

5 files changed, 80 insertions, 19 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/DeviceList.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/DeviceList.kt
new file mode 100644
index 0000000000..de1acb0e23
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/DeviceList.kt
@@ -0,0 +1,7 @@
+package net.mullvad.mullvadvpn.model
+
+sealed class DeviceList {
+    object Unavailable : DeviceList()
+    data class Available(val devices: List<Device>) : DeviceList()
+    object Error : DeviceList()
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
index 42c296159c..2ddc7620b3 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
@@ -11,22 +11,20 @@ import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.withTimeoutOrNull
 import net.mullvad.mullvadvpn.ipc.Event
-import net.mullvad.mullvadvpn.model.Device
+import net.mullvad.mullvadvpn.model.DeviceList
 import net.mullvad.mullvadvpn.model.DeviceListEvent
 import net.mullvad.mullvadvpn.model.DeviceState
 
 class DeviceRepository(
     private val serviceConnectionManager: ServiceConnectionManager,
-    private val deviceListTimeoutMillis: Long = 5000L,
     dispatcher: CoroutineDispatcher = Dispatchers.IO
 ) {
-    private val cachedDeviceList = MutableStateFlow<List<Device>>(emptyList())
+    private val cachedDeviceList = MutableStateFlow<DeviceList>(DeviceList.Unavailable)
 
     val deviceState = serviceConnectionManager.connectionState
         .flatMapLatest { state ->
@@ -55,13 +53,22 @@ class DeviceRepository(
         }
 
     val deviceList = deviceListEvents
-        .map { (it as? DeviceListEvent.Available)?.devices ?: emptyList() }
+        .map {
+            if (it is DeviceListEvent.Available) {
+                DeviceList.Available(it.devices)
+            } else {
+                DeviceList.Error
+            }
+        }
         .onStart {
-            if (cachedDeviceList.value.isNotEmpty()) {
+            if (cachedDeviceList.value is DeviceList.Available) {
                 emit(cachedDeviceList.value)
             }
         }
-        .stateIn(CoroutineScope(Dispatchers.IO), SharingStarted.WhileSubscribed(), emptyList())
+        .shareIn(
+            CoroutineScope(Dispatchers.IO),
+            SharingStarted.WhileSubscribed()
+        )
 
     val deviceRemovalEvent: SharedFlow<Event.DeviceRemovalEvent> =
         serviceConnectionManager.connectionState
@@ -82,7 +89,7 @@ class DeviceRepository(
     }
 
     fun removeDevice(accountToken: String, deviceId: String) {
-        cachedDeviceList.value = emptyList()
+        clearCache()
         serviceConnectionManager.deviceDataSource()?.removeDevice(accountToken, deviceId)
     }
 
@@ -90,17 +97,44 @@ class DeviceRepository(
         serviceConnectionManager.deviceDataSource()?.refreshDeviceList(accountToken)
     }
 
-    suspend fun getDeviceList(accountToken: String): DeviceListEvent {
-        return withTimeoutOrNull(deviceListTimeoutMillis) {
+    fun clearCache() {
+        cachedDeviceList.value = DeviceList.Unavailable
+    }
+
+    private fun updateCache(event: DeviceListEvent, accountToken: String) {
+        cachedDeviceList.value =
+            if (event is DeviceListEvent.Available && event.accountToken == accountToken) {
+                DeviceList.Available(event.devices)
+            } else if (event is DeviceListEvent.Error) {
+                DeviceList.Error
+            } else {
+                DeviceList.Unavailable
+            }
+    }
+
+    suspend fun refreshAndAwaitDeviceListWithTimeout(
+        accountToken: String,
+        shouldClearCache: Boolean,
+        shouldOverrideCache: Boolean,
+        timeoutMillis: Long,
+    ): DeviceListEvent {
+        if (shouldClearCache) {
+            clearCache()
+        }
+
+        val result = withTimeoutOrNull(timeoutMillis) {
             deviceListEvents
                 .onStart {
                     refreshDeviceList(accountToken)
                 }
-                .onEach {
-                    cachedDeviceList.value =
-                        (it as? DeviceListEvent.Available)?.devices ?: emptyList()
-                }
                 .firstOrNull() ?: DeviceListEvent.Error
         } ?: DeviceListEvent.Error
+
+        if (shouldOverrideCache) {
+            updateCache(result, accountToken)
+        }
+
+        return result
     }
+
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
index 5631f74b73..7617310e85 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
@@ -10,6 +10,7 @@ import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.stateIn
 import net.mullvad.mullvadvpn.compose.state.DeviceListItemUiState
 import net.mullvad.mullvadvpn.compose.state.DeviceListUiState
+import net.mullvad.mullvadvpn.model.DeviceList
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.util.safeLet
 
@@ -27,12 +28,17 @@ class DeviceListViewModel(
 
     val uiState = deviceRepository.deviceList
         .combine(_stagedDeviceId) { deviceList, stagedDeviceId ->
-            val stagedDevice = deviceList.firstOrNull { device ->
+            val devices = (deviceList as? DeviceList.Available)?.devices
+            val deviceUiItems = devices?.map {
+                DeviceListItemUiState(it, false)
+            } ?: emptyList()
+            val isLoading = devices == null
+            val stagedDevice = devices?.firstOrNull { device ->
                 device.id == stagedDeviceId
             }
             DeviceListUiState(
-                deviceUiItems = deviceList.map { DeviceListItemUiState(it, false) },
-                isLoading = false,
+                deviceUiItems = deviceUiItems,
+                isLoading = isLoading,
                 stagedDevice = stagedDevice
             )
         }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
index 8751089cf8..e73d42438a 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
@@ -80,7 +80,14 @@ class LoginViewModel(
             LoginResult.Ok -> LoginUiState.Success(false)
             LoginResult.InvalidAccount -> LoginUiState.InvalidAccountError
             LoginResult.MaxDevicesReached -> {
-                if (deviceRepository.getDeviceList(accountToken).isAvailable()) {
+                val refreshResult = deviceRepository.refreshAndAwaitDeviceListWithTimeout(
+                    accountToken = accountToken,
+                    shouldClearCache = true,
+                    shouldOverrideCache = true,
+                    timeoutMillis = 5000L
+                )
+
+                if (refreshResult.isAvailable()) {
                     LoginUiState.TooManyDevicesError(accountToken)
                 } else {
                     LoginUiState.TooManyDevicesMissingListError
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
index 4af15c0cee..c6ae770969 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
@@ -108,7 +108,14 @@ class LoginViewModelTest {
 
     @Test
     fun testLoginWithTooManyDevicesError() = runBlockingTest {
-        coEvery { mockedDeviceRepository.getDeviceList(any()) } returns DeviceListEvent.Available(
+        coEvery {
+            mockedDeviceRepository.refreshAndAwaitDeviceListWithTimeout(
+                any(),
+                any(),
+                any(),
+                any()
+            )
+        } returns DeviceListEvent.Available(
             DUMMY_ACCOUNT_TOKEN, listOf()
         )