Merge branch 'improve-location-info-cache'

9 files changed, 89 insertions, 50 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/ConnectionProxy.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ConnectionProxy.kt
index e55c696a2b..2b551538e2 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/ConnectionProxy.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ConnectionProxy.kt
@@ -1,4 +1,4 @@
-package net.mullvad.mullvadvpn.dataproxy
+package net.mullvad.mullvadvpn.service
 
 import android.content.Context
 import android.content.Intent
@@ -10,7 +10,6 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.model.TunnelState
-import net.mullvad.mullvadvpn.service.MullvadDaemon
 import net.mullvad.mullvadvpn.ui.MainActivity
 import net.mullvad.talpid.tunnel.ActionAfterDisconnect
 import net.mullvad.talpid.util.EventNotifier
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ForegroundNotificationManager.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ForegroundNotificationManager.kt
index 8490cb0d50..97a7e6b99e 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ForegroundNotificationManager.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ForegroundNotificationManager.kt
@@ -12,7 +12,6 @@ import android.content.IntentFilter
 import android.os.Build
 import android.support.v4.app.NotificationCompat
 import net.mullvad.mullvadvpn.R
-import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.mullvadvpn.model.TunnelState
 import net.mullvad.mullvadvpn.ui.MainActivity
 import net.mullvad.talpid.tunnel.ActionAfterDisconnect
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/LocationInfoCache.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/LocationInfoCache.kt
index a521bcebfd..6d6c1e6359 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/dataproxy/LocationInfoCache.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/LocationInfoCache.kt
@@ -1,4 +1,4 @@
-package net.mullvad.mullvadvpn.dataproxy
+package net.mullvad.mullvadvpn.service
 
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.GlobalScope
@@ -11,7 +11,7 @@ import net.mullvad.mullvadvpn.model.TunnelState
 import net.mullvad.mullvadvpn.relaylist.Relay
 import net.mullvad.mullvadvpn.relaylist.RelayCity
 import net.mullvad.mullvadvpn.relaylist.RelayCountry
-import net.mullvad.mullvadvpn.service.MullvadDaemon
+import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.talpid.ConnectivityListener
 import net.mullvad.talpid.tunnel.ActionAfterDisconnect
 
@@ -21,19 +21,27 @@ const val MAX_RETRIES: Int = 17 // ceil(log2(MAX_DELAY / DELAY_SCALE) + 1)
 
 class LocationInfoCache(
     val daemon: MullvadDaemon,
-    val connectivityListener: ConnectivityListener,
-    val relayListListener: RelayListListener
+    val connectionProxy: ConnectionProxy,
+    val connectivityListener: ConnectivityListener
 ) {
-    private var lastKnownRealLocation: GeoIpLocation? = null
     private var activeFetch: Job? = null
+    private var lastKnownRealLocation: GeoIpLocation? = null
+    private var selectedRelayLocation: GeoIpLocation? = null
+
+    private var fetchIdCounter = 0L
+    private var fetchIdIsActive = false
 
     private val connectivityListenerId =
         connectivityListener.connectivityNotifier.subscribe { isConnected ->
-            if (isConnected) {
-                fetchLocation()
+            if (isConnected && state is TunnelState.Disconnected) {
+                fetchLocation(true)
             }
         }
 
+    private val realStateListenerId = connectionProxy.onStateChange.subscribe { realState ->
+        state = realState
+    }
+
     var onNewLocation: ((GeoIpLocation?) -> Unit)? = null
         set(value) {
             field = value
@@ -49,60 +57,88 @@ class LocationInfoCache(
     var state: TunnelState = TunnelState.Disconnected()
         set(value) {
             field = value
+            cancelFetch()
 
             when (value) {
                 is TunnelState.Disconnected -> {
                     location = lastKnownRealLocation
-                    fetchLocation()
+                    fetchLocation(true)
                 }
                 is TunnelState.Connecting -> location = value.location
                 is TunnelState.Connected -> {
                     location = value.location
-                    fetchLocation()
+                    fetchLocation(false)
                 }
                 is TunnelState.Disconnecting -> {
                     when (value.actionAfterDisconnect) {
                         ActionAfterDisconnect.Nothing -> location = lastKnownRealLocation
                         ActionAfterDisconnect.Block -> location = null
-                        ActionAfterDisconnect.Reconnect -> location = locationFromSelectedRelay()
+                        ActionAfterDisconnect.Reconnect -> location = selectedRelayLocation
                     }
                 }
                 is TunnelState.Error -> location = null
             }
         }
 
+    var selectedRelay: RelayItem? = null
+        set(value) {
+            if (field != value) {
+                field = value
+                updateSelectedRelayLocation(value)
+            }
+        }
+
     fun onDestroy() {
         connectivityListener.connectivityNotifier.unsubscribe(connectivityListenerId)
-        activeFetch?.cancel()
+        connectionProxy.onStateChange.unsubscribe(realStateListenerId)
+        cancelFetch()
     }
 
-    private fun locationFromSelectedRelay(): GeoIpLocation? {
-        val relayItem = relayListListener.selectedRelayItem
-
-        when (relayItem) {
-            is RelayCountry -> return GeoIpLocation(null, null, relayItem.name, null, null)
-            is RelayCity -> return GeoIpLocation(
+    private fun updateSelectedRelayLocation(relayItem: RelayItem?) {
+        selectedRelayLocation = when (relayItem) {
+            is RelayCountry -> GeoIpLocation(null, null, relayItem.name, null, null)
+            is RelayCity -> GeoIpLocation(
                 null,
                 null,
                 relayItem.country.name,
                 relayItem.name,
                 null
             )
-            is Relay -> return GeoIpLocation(
+            is Relay -> GeoIpLocation(
                 null,
                 null,
                 relayItem.city.country.name,
                 relayItem.city.name,
                 relayItem.name
             )
+            else -> null
+        }
+    }
+
+    private fun newFetchId(): Long {
+        synchronized(this) {
+            if (fetchIdIsActive) {
+                fetchIdCounter += 1
+            } else {
+                fetchIdIsActive = true
+            }
+
+            return fetchIdCounter
         }
+    }
 
-        return null
+    private fun cancelFetch() {
+        synchronized(this) {
+            if (fetchIdIsActive) {
+                fetchIdCounter += 1
+                fetchIdIsActive = false
+            }
+        }
     }
 
-    private fun fetchLocation() {
+    private fun fetchLocation(isRealLocation: Boolean) {
+        val fetchId = newFetchId()
         val previousFetch = activeFetch
-        val initialState = state
 
         activeFetch = GlobalScope.launch(Dispatchers.Main) {
             var newLocation: GeoIpLocation? = null
@@ -110,21 +146,20 @@ class LocationInfoCache(
 
             previousFetch?.join()
 
-            while (newLocation == null && shouldRetryFetch() && state == initialState) {
+            while (newLocation == null && fetchId == fetchIdCounter) {
                 delayFetch(retry)
                 retry += 1
 
                 newLocation = executeFetch().await()
             }
 
-            if (newLocation != null && state == initialState) {
-                when (state) {
-                    is TunnelState.Disconnected -> {
+            synchronized(this@LocationInfoCache) {
+                if (newLocation != null && fetchId == fetchIdCounter) {
+                    location = newLocation
+
+                    if (isRealLocation) {
                         lastKnownRealLocation = newLocation
-                        location = newLocation
                     }
-                    is TunnelState.Connecting -> location = newLocation
-                    is TunnelState.Connected -> location = newLocation
                 }
             }
         }
@@ -147,11 +182,4 @@ class LocationInfoCache(
 
         delay(duration)
     }
-
-    private fun shouldRetryFetch(): Boolean {
-        val state = this.state
-
-        return connectivityListener.isConnected &&
-            (state is TunnelState.Disconnected || state is TunnelState.Connected)
-    }
 }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadVpnService.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadVpnService.kt
index 1739f77797..4df3ff6791 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadVpnService.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadVpnService.kt
@@ -8,7 +8,6 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.GlobalScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
-import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.mullvadvpn.service.tunnelstate.TunnelStateUpdater
 import net.mullvad.mullvadvpn.ui.MainActivity
 import net.mullvad.talpid.TalpidVpnService
@@ -33,6 +32,7 @@ class MullvadVpnService : TalpidVpnService() {
 
     private var connectionProxy: ConnectionProxy? = null
     private var daemon: MullvadDaemon? = null
+    private var locationInfoCache: LocationInfoCache? = null
     private var startDaemonJob: Job? = null
 
     private lateinit var notificationManager: ForegroundNotificationManager
@@ -156,6 +156,7 @@ class MullvadVpnService : TalpidVpnService() {
             }
 
             onDaemonStopped = {
+                locationInfoCache?.onDestroy()
                 connectionProxy?.onDestroy()
                 serviceNotifier.notify(null)
 
@@ -175,10 +176,19 @@ class MullvadVpnService : TalpidVpnService() {
             pendingAction = null
         }
 
+        val newLocationInfoCache =
+            LocationInfoCache(newDaemon, newConnectionProxy, connectivityListener)
+
         daemon = newDaemon
         connectionProxy = newConnectionProxy
+        locationInfoCache = newLocationInfoCache
 
-        serviceNotifier.notify(ServiceInstance(newDaemon, newConnectionProxy, connectivityListener))
+        serviceNotifier.notify(ServiceInstance(
+            newDaemon,
+            newConnectionProxy,
+            connectivityListener,
+            newLocationInfoCache
+        ))
     }
 
     private fun stop() {
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ServiceInstance.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ServiceInstance.kt
index 070ab1c9bf..8e3e3fbf70 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ServiceInstance.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/ServiceInstance.kt
@@ -1,10 +1,10 @@
 package net.mullvad.mullvadvpn.service
 
-import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.talpid.ConnectivityListener
 
 data class ServiceInstance(
     val daemon: MullvadDaemon,
     val connectionProxy: ConnectionProxy,
-    val connectivityListener: ConnectivityListener
+    val connectivityListener: ConnectivityListener,
+    val locationInfoCache: LocationInfoCache
 )
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/tunnelstate/TunnelStateUpdater.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/tunnelstate/TunnelStateUpdater.kt
index 133c81506e..fa55dbd0cf 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/service/tunnelstate/TunnelStateUpdater.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/service/tunnelstate/TunnelStateUpdater.kt
@@ -1,7 +1,7 @@
 package net.mullvad.mullvadvpn.service.tunnelstate
 
 import android.content.Context
-import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
+import net.mullvad.mullvadvpn.service.ConnectionProxy
 import net.mullvad.mullvadvpn.service.ServiceInstance
 import net.mullvad.talpid.util.EventNotifier
 
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ConnectFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ConnectFragment.kt
index 6c8bfc442d..c59f1b192d 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ConnectFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ConnectFragment.kt
@@ -24,6 +24,7 @@ class ConnectFragment : ServiceDependentFragment(OnNoService.GoToLaunchScreen) {
     private lateinit var locationInfo: LocationInfo
 
     private lateinit var updateKeyStatusJob: Job
+    private var updateLocationInfoJob: Job? = null
     private var updateTunnelStateJob: Job? = null
 
     private var isTunnelInfoExpanded = false
@@ -82,10 +83,14 @@ class ConnectFragment : ServiceDependentFragment(OnNoService.GoToLaunchScreen) {
         }
 
         locationInfoCache.onNewLocation = { location ->
-            locationInfo.location = location
+            updateLocationInfoJob?.cancel()
+            updateLocationInfoJob = GlobalScope.launch(Dispatchers.Main) {
+                locationInfo.location = location
+            }
         }
 
         relayListListener.onRelayListChange = { _, selectedRelayItem ->
+            locationInfoCache.selectedRelay = selectedRelayItem
             switchLocationButton.location = selectedRelayItem
         }
 
@@ -104,6 +109,7 @@ class ConnectFragment : ServiceDependentFragment(OnNoService.GoToLaunchScreen) {
             connectionProxy.onUiStateChange.unsubscribe(listener)
         }
 
+        updateLocationInfoJob?.cancel()
         updateTunnelStateJob?.cancel()
         notificationBanner.onPause()
 
@@ -122,7 +128,6 @@ class ConnectFragment : ServiceDependentFragment(OnNoService.GoToLaunchScreen) {
     private fun updateTunnelState(uiState: TunnelState, realState: TunnelState) =
         GlobalScope.launch(Dispatchers.Main) {
         notificationBanner.tunnelState = realState
-        locationInfoCache.state = realState
         locationInfo.state = realState
         headerBar.setState(realState)
         status.setState(realState)
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceConnection.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceConnection.kt
index 5250f5d42e..116d581c0e 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceConnection.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceConnection.kt
@@ -3,7 +3,6 @@ package net.mullvad.mullvadvpn.ui
 import net.mullvad.mullvadvpn.dataproxy.AccountCache
 import net.mullvad.mullvadvpn.dataproxy.AppVersionInfoCache
 import net.mullvad.mullvadvpn.dataproxy.KeyStatusListener
-import net.mullvad.mullvadvpn.dataproxy.LocationInfoCache
 import net.mullvad.mullvadvpn.dataproxy.RelayListListener
 import net.mullvad.mullvadvpn.dataproxy.SettingsListener
 import net.mullvad.mullvadvpn.service.ServiceInstance
@@ -12,13 +11,13 @@ class ServiceConnection(private val service: ServiceInstance, val mainActivity:
     val daemon = service.daemon
     val connectionProxy = service.connectionProxy
     val connectivityListener = service.connectivityListener
+    val locationInfoCache = service.locationInfoCache
 
     val keyStatusListener = KeyStatusListener(daemon)
     val settingsListener = SettingsListener(daemon)
     val appVersionInfoCache = AppVersionInfoCache(mainActivity, daemon, settingsListener)
     val accountCache = AccountCache(settingsListener, daemon)
     var relayListListener = RelayListListener(daemon, settingsListener)
-    val locationInfoCache = LocationInfoCache(daemon, connectivityListener, relayListListener)
 
     init {
         appVersionInfoCache.onCreate()
@@ -29,7 +28,6 @@ class ServiceConnection(private val service: ServiceInstance, val mainActivity:
         accountCache.onDestroy()
         appVersionInfoCache.onDestroy()
         keyStatusListener.onDestroy()
-        locationInfoCache.onDestroy()
         relayListListener.onDestroy()
         settingsListener.onDestroy()
     }
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
index 7ee6bd36ee..96b8500820 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
@@ -10,11 +10,11 @@ import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.dataproxy.AccountCache
 import net.mullvad.mullvadvpn.dataproxy.AppVersionInfoCache
-import net.mullvad.mullvadvpn.dataproxy.ConnectionProxy
 import net.mullvad.mullvadvpn.dataproxy.KeyStatusListener
-import net.mullvad.mullvadvpn.dataproxy.LocationInfoCache
 import net.mullvad.mullvadvpn.dataproxy.RelayListListener
 import net.mullvad.mullvadvpn.dataproxy.SettingsListener
+import net.mullvad.mullvadvpn.service.ConnectionProxy
+import net.mullvad.mullvadvpn.service.LocationInfoCache
 import net.mullvad.mullvadvpn.service.MullvadDaemon
 import net.mullvad.talpid.ConnectivityListener