Fix header bar blinking issue

Fix the issue that the header bar in ConnectFragment shows
disconnected state very briefly when coming back from
switch location and settings views

Move all states connected to ServiceConnectionManager
to the viewModel to preserve them independently
of the view lifecycle.

Also merge all service connection states to a single
ui state using flow combine.

5 files changed, 150 insertions, 88 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/ConnectUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/ConnectUiState.kt
new file mode 100644
index 0000000000..a69f96a891
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/ConnectUiState.kt
@@ -0,0 +1,27 @@
+package net.mullvad.mullvadvpn.compose.state
+
+import net.mullvad.mullvadvpn.model.GeoIpLocation
+import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.relaylist.RelayItem
+import net.mullvad.mullvadvpn.ui.VersionInfo
+
+data class ConnectUiState(
+    val location: GeoIpLocation?,
+    val relayLocation: RelayItem?,
+    val versionInfo: VersionInfo?,
+    val tunnelUiState: TunnelState,
+    val tunnelRealState: TunnelState,
+    val isTunnelInfoExpanded: Boolean
+) {
+    companion object {
+        val INITIAL =
+            ConnectUiState(
+                location = null,
+                relayLocation = null,
+                versionInfo = null,
+                tunnelUiState = TunnelState.Disconnected,
+                tunnelRealState = TunnelState.Disconnected,
+                isTunnelInfoExpanded = false
+            )
+    }
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
index 1d1f02138e..2a5f82b7aa 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
@@ -80,7 +80,7 @@ val uiModule = module {
     single<IChangelogDataProvider> { ChangelogDataProvider(get()) }
 
     // View models
-    viewModel { ConnectViewModel() }
+    viewModel { ConnectViewModel(get()) }
     viewModel { DeviceRevokedViewModel(get(), get()) }
     viewModel { DeviceListViewModel(get(), get()) }
     viewModel { LoginViewModel(get(), get()) }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragment/ConnectFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragment/ConnectFragment.kt
index 598d460358..803638caa8 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragment/ConnectFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragment/ConnectFragment.kt
@@ -11,15 +11,6 @@ import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.channels.awaitClose
-import kotlinx.coroutines.flow.SharingStarted
-import kotlinx.coroutines.flow.callbackFlow
-import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.debounce
-import kotlinx.coroutines.flow.emptyFlow
-import kotlinx.coroutines.flow.flatMapLatest
-import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.TunnelState
@@ -33,18 +24,13 @@ import net.mullvad.mullvadvpn.ui.notification.AccountExpiryNotification
 import net.mullvad.mullvadvpn.ui.notification.TunnelStateNotification
 import net.mullvad.mullvadvpn.ui.notification.VersionInfoNotification
 import net.mullvad.mullvadvpn.ui.paintNavigationBar
-import net.mullvad.mullvadvpn.ui.serviceconnection.LocationInfoCache
-import net.mullvad.mullvadvpn.ui.serviceconnection.RelayListListener
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
-import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
 import net.mullvad.mullvadvpn.ui.serviceconnection.authTokenCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.connectionProxy
 import net.mullvad.mullvadvpn.ui.widget.HeaderBar
 import net.mullvad.mullvadvpn.ui.widget.NotificationBanner
 import net.mullvad.mullvadvpn.ui.widget.SwitchLocationButton
 import net.mullvad.mullvadvpn.util.JobTracker
-import net.mullvad.mullvadvpn.util.appVersionCallbackFlow
-import net.mullvad.mullvadvpn.util.callbackFlowFromNotifier
 import net.mullvad.mullvadvpn.viewmodel.ConnectViewModel
 import net.mullvad.talpid.tunnel.ErrorStateCause
 import org.koin.android.ext.android.inject
@@ -83,7 +69,7 @@ class ConnectFragment : BaseFragment(), NavigationBarPainter {
 
         headerBar =
             view.findViewById<HeaderBar>(R.id.header_bar).apply {
-                tunnelState = TunnelState.Disconnected
+                tunnelState = connectViewModel.uiState.value.tunnelUiState
             }
 
         accountExpiryNotification.onClick = {
@@ -136,72 +122,24 @@ class ConnectFragment : BaseFragment(), NavigationBarPainter {
         paintNavigationBar(ContextCompat.getColor(requireContext(), R.color.blue))
     }
 
-    val shared =
-        serviceConnectionManager.connectionState
-            .flatMapLatest { state ->
-                if (state is ServiceConnectionState.ConnectedReady) {
-                    flowOf(state.container)
-                } else {
-                    emptyFlow()
-                }
-            }
-            .shareIn(lifecycleScope, SharingStarted.WhileSubscribed())
-
     private fun CoroutineScope.launchUiSubscriptionsOnResume() = launch {
         repeatOnLifecycle(Lifecycle.State.RESUMED) {
-            launchLocationSubscription()
-            launchRelayLocationSubscription()
-            launchTunnelStateSubscription()
-            launchVersionInfoSubscription()
+            launchViewModelSubscription()
             launchAccountExpirySubscription()
-            launchTunnelInfoExpansionSubscription()
         }
     }
 
-    private fun CoroutineScope.launchLocationSubscription() = launch {
-        shared
-            .flatMapLatest { it.locationInfoCache.locationCallbackFlow() }
-            .collect { locationInfo.location = it }
-    }
-
-    private fun LocationInfoCache.locationCallbackFlow() = callbackFlow {
-        onNewLocation = { this.trySend(it) }
-        awaitClose { onNewLocation = null }
-    }
-
-    private fun CoroutineScope.launchRelayLocationSubscription() = launch {
-        shared
-            .flatMapLatest { it.relayListListener.relayListCallbackFlow() }
-            .collect { switchLocationButton.location = it }
-    }
-
-    private fun RelayListListener.relayListCallbackFlow() = callbackFlow {
-        onRelayListChange = { _, item -> this.trySend(item) }
-        awaitClose { onRelayListChange = null }
-    }
-
-    private fun CoroutineScope.launchTunnelStateSubscription() = launch {
-        shared
-            .flatMapLatest {
-                combine(
-                    callbackFlowFromNotifier(it.connectionProxy.onUiStateChange),
-                    callbackFlowFromNotifier(it.connectionProxy.onStateChange)
-                ) { uiState, realState ->
-                    Pair(uiState, realState)
-                }
-            }
-            // Fix to avoid wrong notification shown due to very frequent tunnel state updates.
-            .debounce(TUNNEL_STATE_UPDATE_DEBOUNCE_DURATION_MILLIS)
-            .collect { (uiState, realState) ->
-                tunnelStateNotification.updateTunnelState(uiState)
-                updateTunnelState(uiState, realState)
+    private fun CoroutineScope.launchViewModelSubscription() = launch {
+        connectViewModel.uiState.collect { uiState ->
+            locationInfo.location = uiState.location
+            switchLocationButton.location = uiState.relayLocation
+            uiState.versionInfo?.let {
+                versionInfoNotification.updateVersionInfo(uiState.versionInfo)
             }
-    }
-
-    private fun CoroutineScope.launchVersionInfoSubscription() = launch {
-        shared
-            .flatMapLatest { it.appVersionInfoCache.appVersionCallbackFlow() }
-            .collect { versionInfo -> versionInfoNotification.updateVersionInfo(versionInfo) }
+            tunnelStateNotification.updateTunnelState(uiState.tunnelUiState)
+            updateTunnelState(uiState.tunnelUiState, uiState.tunnelRealState)
+            locationInfo.isTunnelInfoExpanded = uiState.isTunnelInfoExpanded
+        }
     }
 
     private fun CoroutineScope.launchAccountExpirySubscription() = launch {
@@ -210,12 +148,6 @@ class ConnectFragment : BaseFragment(), NavigationBarPainter {
         }
     }
 
-    private fun CoroutineScope.launchTunnelInfoExpansionSubscription() = launch {
-        connectViewModel.isTunnelInfoExpanded.collect { isExpanded ->
-            locationInfo.isTunnelInfoExpanded = isExpanded
-        }
-    }
-
     private fun updateTunnelState(uiState: TunnelState, realState: TunnelState) {
         locationInfo.state = realState
         headerBar.tunnelState = realState
@@ -257,8 +189,4 @@ class ConnectFragment : BaseFragment(), NavigationBarPainter {
             ?.isCausedByExpiredAccount()
             ?: false
     }
-
-    companion object {
-        const val TUNNEL_STATE_UPDATE_DEBOUNCE_DURATION_MILLIS: Long = 200
-    }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
index d7cc4631d0..db742762c6 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
@@ -88,3 +88,26 @@ fun <T> callbackFlowFromNotifier(notifier: EventNotifier<T>) =
         notifier.subscribe(this, handler)
         awaitClose { notifier.unsubscribe(this) }
     }
+
+inline fun <T1, T2, T3, T4, T5, T6, R> combine(
+    flow: Flow<T1>,
+    flow2: Flow<T2>,
+    flow3: Flow<T3>,
+    flow4: Flow<T4>,
+    flow5: Flow<T5>,
+    flow6: Flow<T6>,
+    crossinline transform: suspend (T1, T2, T3, T4, T5, T6) -> R
+): Flow<R> {
+    return kotlinx.coroutines.flow.combine(flow, flow2, flow3, flow4, flow5, flow6) { args: Array<*>
+        ->
+        @Suppress("UNCHECKED_CAST")
+        transform(
+            args[0] as T1,
+            args[1] as T2,
+            args[2] as T3,
+            args[3] as T4,
+            args[4] as T5,
+            args[5] as T6,
+        )
+    }
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModel.kt
index 2be343b6fa..fe8e2662e8 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModel.kt
@@ -1,14 +1,98 @@
 package net.mullvad.mullvadvpn.viewmodel
 
 import androidx.lifecycle.ViewModel
+import androidx.lifecycle.viewModelScope
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.SharedFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.debounce
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.shareIn
+import kotlinx.coroutines.flow.stateIn
+import net.mullvad.mullvadvpn.compose.state.ConnectUiState
+import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.ui.serviceconnection.ConnectionProxy
+import net.mullvad.mullvadvpn.ui.serviceconnection.LocationInfoCache
+import net.mullvad.mullvadvpn.ui.serviceconnection.RelayListListener
+import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
+import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
+import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
+import net.mullvad.mullvadvpn.util.appVersionCallbackFlow
+import net.mullvad.mullvadvpn.util.callbackFlowFromNotifier
+import net.mullvad.mullvadvpn.util.combine
+
+class ConnectViewModel(serviceConnectionManager: ServiceConnectionManager) : ViewModel() {
+    private val _shared: SharedFlow<ServiceConnectionContainer> =
+        serviceConnectionManager.connectionState
+            .flatMapLatest { state ->
+                if (state is ServiceConnectionState.ConnectedReady) {
+                    flowOf(state.container)
+                } else {
+                    emptyFlow()
+                }
+            }
+            .shareIn(viewModelScope, SharingStarted.WhileSubscribed())
 
-class ConnectViewModel : ViewModel() {
     private val _isTunnelInfoExpanded = MutableStateFlow(false)
-    val isTunnelInfoExpanded = _isTunnelInfoExpanded.asStateFlow()
+
+    val uiState: StateFlow<ConnectUiState> =
+        _shared
+            .flatMapLatest { serviceConnection ->
+                combine(
+                    serviceConnection.locationInfoCache.locationCallbackFlow(),
+                    serviceConnection.relayListListener.relayListCallbackFlow(),
+                    serviceConnection.appVersionInfoCache.appVersionCallbackFlow(),
+                    serviceConnection.connectionProxy.tunnelUiStateFlow(),
+                    serviceConnection.connectionProxy.tunnelRealStateFlow(),
+                    _isTunnelInfoExpanded
+                ) {
+                    location,
+                    relayLocation,
+                    versionInfo,
+                    tunnelUiState,
+                    tunnelRealState,
+                    isTunnelInfoExpanded ->
+                    ConnectUiState(
+                        location = location,
+                        relayLocation = relayLocation,
+                        versionInfo = versionInfo,
+                        tunnelUiState = tunnelUiState,
+                        tunnelRealState = tunnelRealState,
+                        isTunnelInfoExpanded = isTunnelInfoExpanded
+                    )
+                }
+            }
+            .stateIn(viewModelScope, SharingStarted.WhileSubscribed(), ConnectUiState.INITIAL)
+
+    private fun LocationInfoCache.locationCallbackFlow() = callbackFlow {
+        onNewLocation = { this.trySend(it) }
+        awaitClose { onNewLocation = null }
+    }
+
+    private fun RelayListListener.relayListCallbackFlow() = callbackFlow {
+        onRelayListChange = { _, item -> this.trySend(item) }
+        awaitClose { onRelayListChange = null }
+    }
+
+    private fun ConnectionProxy.tunnelUiStateFlow(): Flow<TunnelState> =
+        callbackFlowFromNotifier(this.onUiStateChange)
+            .debounce(TUNNEL_STATE_UPDATE_DEBOUNCE_DURATION_MILLIS)
+
+    private fun ConnectionProxy.tunnelRealStateFlow(): Flow<TunnelState> =
+        callbackFlowFromNotifier(this.onStateChange)
+            .debounce(TUNNEL_STATE_UPDATE_DEBOUNCE_DURATION_MILLIS)
 
     fun toggleTunnelInfoExpansion() {
         _isTunnelInfoExpanded.value = _isTunnelInfoExpanded.value.not()
     }
+
+    companion object {
+        const val TUNNEL_STATE_UPDATE_DEBOUNCE_DURATION_MILLIS: Long = 200
+    }
 }