Nodes: rework RootFlowNode with cache service

appnav/src/main/kotlin/io/element/android/appnav/RootFlowNode.kt

@@ -31,7 +31,6 @@ import com.bumble.appyx.core.node.node
 import com.bumble.appyx.core.plugin.Plugin
 import com.bumble.appyx.core.plugin.plugins
 import com.bumble.appyx.navmodel.backstack.BackStack
-import com.bumble.appyx.navmodel.backstack.operation.newRoot
 import com.bumble.appyx.navmodel.backstack.operation.pop
 import com.bumble.appyx.navmodel.backstack.operation.push
 import dagger.assisted.Assisted
@@ -49,19 +48,21 @@ import io.element.android.features.rageshake.api.bugreport.BugReportEntryPoint
 import io.element.android.libraries.architecture.BackstackNode
 import io.element.android.libraries.architecture.animation.rememberDefaultTransitionHandler
 import io.element.android.libraries.architecture.createNode
+import io.element.android.libraries.architecture.waitForChildAttached
 import io.element.android.libraries.deeplink.DeeplinkData
 import io.element.android.libraries.designsystem.theme.components.CircularProgressIndicator
 import io.element.android.libraries.di.AppScope
 import io.element.android.libraries.matrix.api.auth.MatrixAuthenticationService
 import io.element.android.libraries.matrix.api.core.SessionId
-import io.element.android.libraries.matrix.api.core.UserId
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.onStart
 import kotlinx.parcelize.Parcelize
 import timber.log.Timber
+import java.util.UUID
 
 @ContributesNode(AppScope::class)
 class RootFlowNode @AssistedInject constructor(
@@ -90,21 +91,15 @@ class RootFlowNode @AssistedInject constructor(
     }
 
     private fun observeLoggedInState() {
-        authenticationService.isLoggedIn()
+        combine(
-            .distinctUntilChanged()
+            cacheService.onClearedCacheEventFlow(),
-            .combine(
+            authenticationService.isLoggedIn(),
-                cacheService.cacheIndex().onEach {
+        ) { _, isLoggedIn -> isLoggedIn }
-                    Timber.v("cacheIndex=$it")
+            .onEach { isLoggedIn ->
-                    matrixClientsHolder.removeAll()
+                Timber.v("isLoggedIn=$isLoggedIn")
-                }
-            ) { isLoggedIn, cacheIdx -> isLoggedIn to cacheIdx }
-            .onEach { pair ->
-                val isLoggedIn = pair.first
-                val cacheIndex = pair.second
-                Timber.v("isLoggedIn=$isLoggedIn, cacheIndex=$cacheIndex")
                 if (isLoggedIn) {
                     tryToRestoreLatestSession(
-                        onSuccess = { switchToLoggedInFlow(it, cacheIndex) },
+                        onSuccess = { switchToLoggedInFlow(it) },
                         onFailure = { switchToNotLoggedInFlow() }
                     )
                 } else {
@@ -114,8 +109,8 @@ class RootFlowNode @AssistedInject constructor(
             .launchIn(lifecycleScope)
     }
 
-    private fun switchToLoggedInFlow(sessionId: SessionId, cacheIndex: Int) {
+    private fun switchToLoggedInFlow(sessionId: SessionId) {
-        backstack.safeRoot(NavTarget.LoggedInFlow(sessionId, cacheIndex))
+        backstack.safeRoot(NavTarget.LoggedInFlow(sessionId))
     }
 
     private fun switchToNotLoggedInFlow() {
@@ -123,28 +118,38 @@ class RootFlowNode @AssistedInject constructor(
         backstack.safeRoot(NavTarget.NotLoggedInFlow)
     }
 
+    private suspend fun restoreSessionIfNeeded(
+        sessionId: SessionId,
+        onFailure: () -> Unit = {},
+        onSuccess: (SessionId) -> Unit = {},
+    ) {
+        // If the session is already known it'll be restored by the node hierarchy
+        if (matrixClientsHolder.knowSession(sessionId)) {
+            Timber.v("Session $sessionId already alive, no need to restore.")
+            return
+        }
+        authenticationService.restoreSession(sessionId)
+            .onSuccess { matrixClient ->
+                matrixClientsHolder.add(matrixClient)
+                Timber.v("Succeed to restore session $sessionId")
+                onSuccess(sessionId)
+            }
+            .onFailure {
+                Timber.v("Failed to restore session $sessionId")
+                onFailure()
+            }
+    }
+
     private suspend fun tryToRestoreLatestSession(
-        onSuccess: (UserId) -> Unit = {},
+        onSuccess: (SessionId) -> Unit = {},
         onFailure: () -> Unit = {}
     ) {
-        val latestKnownUserId = authenticationService.getLatestSessionId()
+        val latestSessionId = authenticationService.getLatestSessionId()
-        if (latestKnownUserId == null) {
+        if (latestSessionId == null) {
             onFailure()
             return
         }
-        if (matrixClientsHolder.knowSession(latestKnownUserId)) {
+        restoreSessionIfNeeded(latestSessionId, onFailure, onSuccess)
-            onSuccess(latestKnownUserId)
-            return
-        }
-        authenticationService.restoreSession(UserId(latestKnownUserId.value))
-            .onSuccess { matrixClient ->
-                matrixClientsHolder.add(matrixClient)
-                onSuccess(matrixClient.sessionId)
-            }
-            .onFailure {
-                Timber.v("Failed to restore session...")
-                onFailure()
-            }
     }
 
     private fun onOpenBugReport() {
@@ -175,7 +180,10 @@ class RootFlowNode @AssistedInject constructor(
         object NotLoggedInFlow : NavTarget
 
         @Parcelize
-        data class LoggedInFlow(val sessionId: SessionId, val cacheIndex: Int) : NavTarget
+        data class LoggedInFlow(
+            val sessionId: SessionId,
+            val navId: UUID = UUID.randomUUID(),
+        ) : NavTarget
 
         @Parcelize
         object BugReport : NavTarget
@@ -186,7 +194,6 @@ class RootFlowNode @AssistedInject constructor(
             is NavTarget.LoggedInFlow -> {
                 val matrixClient = matrixClientsHolder.getOrNull(navTarget.sessionId) ?: return splashNode(buildContext).also {
                     Timber.w("Couldn't find any session, go through SplashScreen")
-                    backstack.newRoot(NavTarget.SplashScreen)
                 }
                 val inputs = LoggedInFlowNode.Inputs(matrixClient)
                 val callback = object : LoggedInFlowNode.Callback {
@@ -247,9 +254,16 @@ class RootFlowNode @AssistedInject constructor(
     }
 
     private suspend fun attachSession(sessionId: SessionId): LoggedInFlowNode {
-        val cacheIndex = cacheService.cacheIndex().first()
+        //TODO handle multi-session
-        return attachChild {
+        return waitForChildAttached { navTarget ->
-            backstack.newRoot(NavTarget.LoggedInFlow(sessionId, cacheIndex))
+            navTarget is NavTarget.LoggedInFlow && navTarget.sessionId == sessionId
         }
     }
+
+    private fun CacheService.onClearedCacheEventFlow(): Flow<Unit> {
+        return clearedCacheEventFlow
+            .onEach { sessionId -> matrixClientsHolder.remove(sessionId) }
+            .map { }
+            .onStart { emit((Unit)) }
+    }
 }

features/preferences/api/build.gradle.kts

@@ -23,4 +23,5 @@ android {
 
 dependencies {
     implementation(projects.libraries.architecture)
+    api(projects.libraries.matrix.api)
 }

features/preferences/api/src/main/kotlin/io/element/android/features/preferences/api/CacheService.kt

@@ -16,13 +16,13 @@
 
 package io.element.android.features.preferences.api
 
+import io.element.android.libraries.matrix.api.core.SessionId
 import kotlinx.coroutines.flow.Flow
 
 interface CacheService {
     /**
-     * Returns a flow of the current cache index, can let the app to know when the
+     * A flow of [SessionId], can let the app to know when the
-     * cache has been cleared, for instance to restart the app.
+     * cache has been cleared for a given session, for instance to restart the app.
-     * Will be a flow of Int, starting from 0, and incrementing each time the cache is cleared.
      */
-    fun cacheIndex(): Flow<Int>
+    val clearedCacheEventFlow: Flow<SessionId>
 }

features/preferences/impl/src/main/kotlin/io/element/android/features/preferences/impl/DefaultCacheService.kt

@@ -20,20 +20,19 @@ import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.features.preferences.api.CacheService
 import io.element.android.libraries.di.AppScope
 import io.element.android.libraries.di.SingleIn
+import io.element.android.libraries.matrix.api.core.SessionId
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import javax.inject.Inject
 
 @SingleIn(AppScope::class)
 @ContributesBinding(AppScope::class)
 class DefaultCacheService @Inject constructor() : CacheService {
-    private val cacheIndexState = MutableStateFlow(0)
 
-    override fun cacheIndex(): Flow<Int> {
+    private val _clearedCacheEventFlow = MutableSharedFlow<SessionId>(0)
-        return cacheIndexState
+    override val clearedCacheEventFlow: Flow<SessionId> = _clearedCacheEventFlow
-    }
 
-    fun incrementCacheIndex() {
+    suspend fun onClearedCache(sessionId: SessionId) {
-        cacheIndexState.value++
+        _clearedCacheEventFlow.emit(sessionId)
     }
 }

features/preferences/impl/src/main/kotlin/io/element/android/features/preferences/impl/tasks/ClearCacheUseCase.kt

@@ -27,6 +27,7 @@ import io.element.android.libraries.core.coroutine.CoroutineDispatchers
 import io.element.android.libraries.di.ApplicationContext
 import io.element.android.libraries.di.SessionScope
 import io.element.android.libraries.matrix.api.MatrixClient
+import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.withContext
 import okhttp3.OkHttpClient
 import javax.inject.Inject
@@ -57,6 +58,6 @@ class DefaultClearCacheUseCase @Inject constructor(
         // Clear app cache
         context.cacheDir.deleteRecursively()
         // Ensure the app is restarted
-        defaultCacheIndexProvider.incrementCacheIndex()
+        defaultCacheIndexProvider.onClearedCache(matrixClient.sessionId)
     }
 }

libraries/architecture/src/main/kotlin/io/element/android/libraries/architecture/ParentNodeExt.kt

@@ -16,12 +16,35 @@
 
 package io.element.android.libraries.architecture
 
+import androidx.lifecycle.lifecycleScope
 import com.bumble.appyx.core.children.nodeOrNull
 import com.bumble.appyx.core.node.Node
 import com.bumble.appyx.core.node.ParentNode
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlin.coroutines.resume
 
 fun <NavTarget : Any> ParentNode<NavTarget>.childNode(navTarget: NavTarget): Node? {
     val childMap = children.value
     val key = childMap.keys.find { it.navTarget == navTarget }
     return childMap[key]?.nodeOrNull
 }
+
+suspend inline fun <reified N : Node, NavTarget : Any> ParentNode<NavTarget>.waitForChildAttached(crossinline predicate: (NavTarget) -> Boolean): N =
+    suspendCancellableCoroutine { continuation ->
+        lifecycleScope.launch {
+            children.collect { childMap ->
+                val expectedChildNode = childMap.entries
+                    .map { it.key.navTarget }
+                    .lastOrNull(predicate)
+                    ?.let {
+                        childNode(it) as? N
+                    }
+                if (expectedChildNode != null && !continuation.isCompleted) {
+                    continuation.resume(expectedChildNode)
+                }
+            }
+        }.invokeOnCompletion {
+            continuation.cancel()
+        }
+    }