Handle activity and process recreation for session

app/src/main/kotlin/io/element/android/x/node/BackstackExt.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2023 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.element.android.x.node
+
+import com.bumble.appyx.navmodel.backstack.BackStack
+import com.bumble.appyx.navmodel.backstack.operation.NewRoot
+
+/**
+ * Don't process NewRoot if the nav target already exists in the stack.
+ */
+fun <T : Any> BackStack<T>.safeRoot(element: T) {
+    val containsRoot = elements.value.any {
+        it.key.navTarget == element
+    }
+    if (containsRoot) return
+    accept(NewRoot(element))
+}

app/src/main/kotlin/io/element/android/x/node/RootFlowNode.kt

@@ -27,7 +27,6 @@ import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.lifecycle.lifecycleScope
 import com.bumble.appyx.core.composable.Children
-import com.bumble.appyx.core.lifecycle.subscribe
 import com.bumble.appyx.core.modality.BuildContext
 import com.bumble.appyx.core.node.Node
 import com.bumble.appyx.core.node.ParentNode
@@ -41,9 +40,9 @@ import io.element.android.libraries.architecture.animation.rememberDefaultTransi
 import io.element.android.libraries.architecture.createNode
 import io.element.android.libraries.architecture.presenterConnector
 import io.element.android.libraries.di.DaggerComponentOwner
-import io.element.android.libraries.matrix.MatrixClient
 import io.element.android.libraries.matrix.auth.MatrixAuthenticationService
 import io.element.android.libraries.matrix.core.SessionId
+import io.element.android.x.di.MatrixClientsHolder
 import io.element.android.x.root.RootPresenter
 import io.element.android.x.root.RootView
 import kotlinx.coroutines.flow.distinctUntilChanged
@@ -51,54 +50,81 @@ import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.parcelize.Parcelize
 import timber.log.Timber
-import java.util.concurrent.ConcurrentHashMap
 
 class RootFlowNode(
-    buildContext: BuildContext,
+    private val buildContext: BuildContext,
     private val backstack: BackStack<NavTarget> = BackStack(
         initialElement = NavTarget.SplashScreen,
         savedStateMap = buildContext.savedStateMap,
     ),
     private val appComponentOwner: DaggerComponentOwner,
     private val authenticationService: MatrixAuthenticationService,
+    private val matrixClientsHolder: MatrixClientsHolder,
     presenter: RootPresenter
 ) :
     ParentNode<RootFlowNode.NavTarget>(
         navModel = backstack,
-        buildContext = buildContext,
+        buildContext = buildContext
     ),
 
     DaggerComponentOwner by appComponentOwner {
 
-    private val matrixClientsHolder = ConcurrentHashMap<SessionId, MatrixClient>()
     private val presenterConnector = presenterConnector(presenter)
 
     override fun onBuilt() {
         super.onBuilt()
-        whenChildAttached(LoggedInFlowNode::class) { _, child ->
-            child.lifecycle.subscribe(
-                onDestroy = { matrixClientsHolder.remove(child.sessionId) }
-            )
-        }
+        observeLoggedInState()
+    }
+
+    private fun observeLoggedInState() {
         authenticationService.isLoggedIn()
             .distinctUntilChanged()
             .onEach { isLoggedIn ->
                 Timber.v("isLoggedIn=$isLoggedIn")
                 if (isLoggedIn) {
-                    val matrixClient = authenticationService.restoreSession()
-                    if (matrixClient == null) {
-                        backstack.newRoot(NavTarget.NotLoggedInFlow)
-                    } else {
-                        matrixClientsHolder[matrixClient.sessionId] = matrixClient
-                        backstack.newRoot(NavTarget.LoggedInFlow(matrixClient.sessionId))
-                    }
+                    tryToRestoreLatestSession(
+                        onSuccess = { switchToLoggedInFlow(it) },
+                        onFailure = { switchToLogoutFlow() }
+                    )
                 } else {
-                    backstack.newRoot(NavTarget.NotLoggedInFlow)
+                    switchToLogoutFlow()
                 }
             }
             .launchIn(lifecycleScope)
     }
 
+    private fun switchToLoggedInFlow(sessionId: SessionId) {
+        backstack.safeRoot(NavTarget.LoggedInFlow(sessionId = sessionId))
+    }
+
+    private fun switchToLogoutFlow() {
+        matrixClientsHolder.removeAll()
+        backstack.safeRoot(NavTarget.NotLoggedInFlow)
+    }
+
+    private suspend fun tryToRestoreLatestSession(
+        onSuccess: (SessionId) -> Unit = {},
+        onFailure: () -> Unit = {}
+    ) {
+        val latestKnownSessionId = authenticationService.getLatestSessionId()
+        if (latestKnownSessionId == null) {
+            onFailure()
+            return
+        }
+        if (matrixClientsHolder.knowSession(latestKnownSessionId)) {
+            onSuccess(latestKnownSessionId)
+            return
+        }
+        val matrixClient = authenticationService.restoreSession(latestKnownSessionId)
+        if (matrixClient == null) {
+            Timber.v("Failed to restore session...")
+            onFailure()
+        } else {
+            matrixClientsHolder.add(matrixClient)
+            onSuccess(matrixClient.sessionId)
+        }
+    }
+
     private fun onOpenBugReport() {
         backstack.push(NavTarget.BugReport)
     }
@@ -142,8 +168,10 @@ class RootFlowNode(
     override fun resolve(navTarget: NavTarget, buildContext: BuildContext): Node {
         return when (navTarget) {
             is NavTarget.LoggedInFlow -> {
-                val matrixClient =
-                    matrixClientsHolder[navTarget.sessionId] ?: throw IllegalStateException("Makes sure to give a matrixClient with the given sessionId")
+                val matrixClient = matrixClientsHolder.getOrNull(navTarget.sessionId) ?: return splashNode(buildContext).also {
+                    Timber.w("Couldn't find any session, go through SplashScreen")
+                    backstack.newRoot(NavTarget.SplashScreen)
+                }
                 LoggedInFlowNode(
                     buildContext = buildContext,
                     sessionId = navTarget.sessionId,
@@ -152,12 +180,14 @@ class RootFlowNode(
                 )
             }
             NavTarget.NotLoggedInFlow -> NotLoggedInFlowNode(buildContext)
-            NavTarget.SplashScreen -> node(buildContext) {
-                Box(modifier = it.fillMaxSize(), contentAlignment = Alignment.Center) {
-                    CircularProgressIndicator()
-                }
-            }
+            NavTarget.SplashScreen -> splashNode(buildContext)
             NavTarget.BugReport -> createNode<BugReportNode>(buildContext, plugins = listOf(bugReportNodeCallback))
         }
     }
+
+    private fun splashNode(buildContext: BuildContext) = node(buildContext) {
+        Box(modifier = it.fillMaxSize(), contentAlignment = Alignment.Center) {
+            CircularProgressIndicator()
+        }
+    }
 }