Handle activity and process recreation for session

app/src/main/kotlin/io/element/android/x/MainActivity.kt

@@ -36,6 +36,7 @@ class MainActivity : NodeComponentActivity() {
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
         val appBindings = bindings<AppBindings>()
+        appBindings.matrixClientsHolder().restore(savedInstanceState)
         WindowCompat.setDecorFitsSystemWindows(window, false)
         setContent {
             ElementXTheme {
@@ -48,11 +49,17 @@ class MainActivity : NodeComponentActivity() {
                             buildContext = it,
                             appComponentOwner = applicationContext as DaggerComponentOwner,
                             authenticationService = appBindings.authenticationService(),
-                            presenter = appBindings.rootPresenter()
+                            presenter = appBindings.rootPresenter(),
+                            matrixClientsHolder = appBindings.matrixClientsHolder()
                         )
                     }
                 }
             }
         }
     }
+
+    override fun onSaveInstanceState(outState: Bundle) {
+        super.onSaveInstanceState(outState)
+        bindings<AppBindings>().matrixClientsHolder().onSaveInstanceState(outState)
+    }
 }

app/src/main/kotlin/io/element/android/x/di/AppBindings.kt

@@ -24,7 +24,7 @@ import kotlinx.coroutines.CoroutineScope
 
 @ContributesTo(AppScope::class)
 interface AppBindings {
-    fun coroutineScope(): CoroutineScope
     fun rootPresenter(): RootPresenter
     fun authenticationService(): MatrixAuthenticationService
+    fun matrixClientsHolder(): MatrixClientsHolder
 }

app/src/main/kotlin/io/element/android/x/di/MatrixClientsHolder.kt

@@ -0,0 +1,79 @@
+/*
+ * 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.di
+
+import android.os.Bundle
+import io.element.android.libraries.di.AppScope
+import io.element.android.libraries.di.SingleIn
+import io.element.android.libraries.matrix.MatrixClient
+import io.element.android.libraries.matrix.auth.MatrixAuthenticationService
+import io.element.android.libraries.matrix.core.SessionId
+import kotlinx.coroutines.runBlocking
+import timber.log.Timber
+import java.util.concurrent.ConcurrentHashMap
+import javax.inject.Inject
+
+private const val SAVE_INSTANCE_KEY = "io.element.android.x.di.MatrixClientsHolder.SaveInstanceKey"
+
+@SingleIn(AppScope::class)
+class MatrixClientsHolder @Inject constructor(private val authenticationService: MatrixAuthenticationService) {
+
+    private val sessionIdsToMatrixClient = ConcurrentHashMap<SessionId, MatrixClient>()
+
+    fun add(matrixClient: MatrixClient) {
+        sessionIdsToMatrixClient[matrixClient.sessionId] = matrixClient
+    }
+
+    fun removeAll() {
+        sessionIdsToMatrixClient.clear()
+    }
+
+    fun remove(sessionId: SessionId) {
+        sessionIdsToMatrixClient.remove(sessionId)
+    }
+
+    fun isEmpty(): Boolean = sessionIdsToMatrixClient.isEmpty()
+
+    fun knowSession(sessionId: SessionId): Boolean = sessionIdsToMatrixClient.containsKey(sessionId)
+
+    fun getOrNull(sessionId: SessionId): MatrixClient? {
+        return sessionIdsToMatrixClient[sessionId]
+    }
+
+    @Suppress("DEPRECATION")
+    fun restore(savedInstanceState: Bundle?) {
+        if (savedInstanceState == null || sessionIdsToMatrixClient.isNotEmpty()) return
+        val sessionIds = savedInstanceState.getSerializable(SAVE_INSTANCE_KEY) as? Array<SessionId>
+        if (sessionIds.isNullOrEmpty()) return
+        // Not ideal but should only happens in case of process recreation. This ensure we restore all the active sessions before restoring the node graphs.
+        runBlocking {
+            sessionIds.forEach { sessionId ->
+                Timber.v("Restore matrix session: $sessionId")
+                val matrixClient = authenticationService.restoreSession(sessionId)
+                if (matrixClient != null) {
+                    add(matrixClient)
+                }
+            }
+        }
+    }
+
+    fun onSaveInstanceState(outState: Bundle) {
+        val sessionKeys = sessionIdsToMatrixClient.keys.toTypedArray()
+        Timber.v("Save matrix session keys = $sessionKeys")
+        outState.putSerializable(SAVE_INSTANCE_KEY, sessionKeys)
+    }
+}

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()
+        }
+    }
 }

libraries/matrix/src/main/kotlin/io/element/android/libraries/matrix/auth/MatrixAuthenticationService.kt

@@ -23,7 +23,7 @@ import kotlinx.coroutines.flow.Flow
 interface MatrixAuthenticationService {
     fun isLoggedIn(): Flow<Boolean>
     suspend fun getLatestSessionId(): SessionId?
-    suspend fun restoreSession(): MatrixClient?
+    suspend fun restoreSession(sessionId: SessionId): MatrixClient?
     fun getHomeserver(): String?
     fun getHomeserverOrDefault(): String
     suspend fun setHomeserver(homeserver: String)

libraries/matrix/src/main/kotlin/io/element/android/libraries/matrix/auth/RustMatrixAuthenticationService.kt

@@ -52,8 +52,8 @@ class RustMatrixAuthenticationService @Inject constructor(
         sessionStore.getLatestSession()?.sessionId()
     }
 
-    override suspend fun restoreSession() = withContext(coroutineDispatchers.io) {
-        sessionStore.getLatestSession()
+    override suspend fun restoreSession(sessionId: SessionId) = withContext(coroutineDispatchers.io) {
+        sessionStore.getSession(sessionId)
             ?.let { session ->
                 try {
                     ClientBuilder()

libraries/matrix/src/main/kotlin/io/element/android/libraries/matrix/session/PreferencesSessionStore.kt

@@ -26,6 +26,7 @@ import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.libraries.di.AppScope
 import io.element.android.libraries.di.ApplicationContext
 import io.element.android.libraries.di.SingleIn
+import io.element.android.libraries.matrix.core.SessionId
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.flow.map
@@ -94,6 +95,11 @@ class PreferencesSessionStore @Inject constructor(
         }
     }
 
+    override suspend fun getSession(sessionId: SessionId): Session? {
+        //TODO we should have a proper session management
+        return getLatestSession()
+    }
+
     override suspend fun reset() {
         store.edit { it.clear() }
     }

libraries/matrix/src/main/kotlin/io/element/android/libraries/matrix/session/SessionStore.kt

@@ -16,12 +16,14 @@
 
 package io.element.android.libraries.matrix.session
 
+import io.element.android.libraries.matrix.core.SessionId
 import kotlinx.coroutines.flow.Flow
 import org.matrix.rustcomponents.sdk.Session
 
 interface SessionStore {
     fun isLoggedIn(): Flow<Boolean>
     suspend fun storeData(session: Session)
+    suspend fun getSession(sessionId: SessionId): Session?
     suspend fun getLatestSession(): Session?
     suspend fun reset()
 }