Unregister listener when scope or flow is closed

1 files changed, 28 insertions, 7 deletions

diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/ServiceConnection.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/ServiceConnection.kt
index 6211f79dc7..6683b7459a 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/ServiceConnection.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ipc/ServiceConnection.kt
@@ -7,13 +7,13 @@ import android.os.Looper
 import android.os.Messenger
 import kotlin.reflect.KClass
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.InternalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.service.MullvadVpnService
@@ -21,31 +21,34 @@ import net.mullvad.mullvadvpn.util.DispatchingFlow
 import net.mullvad.mullvadvpn.util.bindServiceFlow
 import net.mullvad.mullvadvpn.util.dispatchTo
 
-@InternalCoroutinesApi
 class ServiceConnection(context: Context, scope: CoroutineScope) {
-    private val connection = MutableStateFlow<Messenger?>(null)
+    private val activeListeners = MutableStateFlow<Pair<Messenger, Int>?>(null)
     private val handler = HandlerFlow(Looper.getMainLooper(), Event::fromMessage)
     private val listener = Messenger(handler)
+    private val listenerId = MutableStateFlow<Int?>(null)
 
-    private lateinit var requestSinks: StateFlow<Messenger?>
+    private lateinit var listenerRegistrations: StateFlow<Pair<Messenger, Int>?>
 
     init {
         val dispatcher = handler
             .filterNotNull()
             .dispatchTo {
-                requestSinks = subscribeToState(Event.ListenerReady::class, scope) { connection }
+                listenerRegistrations = subscribeToState(Event.ListenerReady::class, scope) {
+                    Pair(connection, listenerId)
+                }
             }
 
         scope.launch { connect(context) }
         scope.launch { dispatcher.collect() }
-        scope.launch { requestSinks.collect { connection.value = it } }
+        scope.launch { unregisterOldListeners() }
+        scope.launch { listenerRegistrations.collect { activeListeners.value = it } }
     }
 
     private suspend fun connect(context: Context) {
         val intent = Intent(context, MullvadVpnService::class.java)
 
         context.bindServiceFlow(intent).collect { binder ->
-            connection.value = null
+            activeListeners.value = null
             binder?.let(::registerListener)
         }
     }
@@ -57,6 +60,24 @@ class ServiceConnection(context: Context, scope: CoroutineScope) {
         messenger.send(request.message)
     }
 
+    private suspend fun unregisterOldListeners() {
+        var oldListener: Pair<Messenger, Int>? = null
+
+        activeListeners
+            .onCompletion { oldListener?.let(::unregisterListener) }
+            .collect { newListener ->
+                oldListener?.let(::unregisterListener)
+                oldListener = newListener
+            }
+    }
+
+    private fun unregisterListener(registration: Pair<Messenger, Int>) {
+        val (messenger, listenerId) = registration
+        val request = Request.UnregisterListener(listenerId)
+
+        messenger.send(request.message)
+    }
+
     private fun <V : Any, D> DispatchingFlow<in V>.subscribeToState(
         event: KClass<V>,
         scope: CoroutineScope,