Try to centralise session restoration through MatrixClientsHolder

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

@@ -145,14 +145,10 @@ class RootFlowNode @AssistedInject constructor(
         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
+        runCatching {
+            matrixClientsHolder.requireSession(sessionId)
         }
-        authenticationService.restoreSession(sessionId)
-            .onSuccess { matrixClient ->
-                matrixClientsHolder.add(matrixClient)
+            .onSuccess {
                 Timber.v("Succeed to restore session $sessionId")
                 onSuccess(sessionId)
             }

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

@@ -35,6 +35,6 @@ fun Throwable.mapAuthenticationException(): Throwable {
         is RustAuthenticationException.OidcNotSupported -> AuthenticationException.OidcError("OidcNotSupported", message!!)
          */
 
-        else -> this
+        else -> AuthenticationException.Generic(this.message ?: "Unknown error")
     }
 }

libraries/matrixui/src/main/kotlin/io/element/android/libraries/matrix/ui/di/MatrixClientsHolder.kt

@@ -21,12 +21,16 @@ import com.bumble.appyx.core.state.SavedStateMap
 import io.element.android.libraries.di.AppScope
 import io.element.android.libraries.di.SingleIn
 import io.element.android.libraries.matrix.api.MatrixClient
+import io.element.android.libraries.matrix.api.auth.AuthenticationException
 import io.element.android.libraries.matrix.api.auth.MatrixAuthenticationService
 import io.element.android.libraries.matrix.api.core.SessionId
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 import timber.log.Timber
 import java.util.concurrent.ConcurrentHashMap
 import javax.inject.Inject
+import kotlin.jvm.Throws
 
 private const val SAVE_INSTANCE_KEY = "io.element.android.x.di.MatrixClientsHolder.SaveInstanceKey"
 
@@ -34,8 +38,9 @@ private const val SAVE_INSTANCE_KEY = "io.element.android.x.di.MatrixClientsHold
 class MatrixClientsHolder @Inject constructor(private val authenticationService: MatrixAuthenticationService) {
 
     private val sessionIdsToMatrixClient = ConcurrentHashMap<SessionId, MatrixClient>()
+    private val restoreMutex = Mutex()
 
-    fun add(matrixClient: MatrixClient) {
+    private fun add(matrixClient: MatrixClient) {
         sessionIdsToMatrixClient[matrixClient.sessionId] = matrixClient
     }
 
@@ -55,6 +60,27 @@ class MatrixClientsHolder @Inject constructor(private val authenticationService:
         return sessionIdsToMatrixClient[sessionId]
     }
 
+    @Throws(AuthenticationException::class)
+    suspend fun requireSession(sessionId: SessionId): MatrixClient {
+        return restoreMutex.withLock {
+             when (val matrixClient = sessionIdsToMatrixClient[sessionId]) {
+                null -> restore(sessionId).getOrThrow()
+                else -> matrixClient
+            }
+        }
+    }
+
+    private suspend fun restore(sessionId: SessionId): Result<MatrixClient> {
+        Timber.d("Restore matrix session: $sessionId")
+        return authenticationService.restoreSession(sessionId)
+            .onSuccess { matrixClient ->
+                add(matrixClient)
+            }
+            .onFailure {
+                Timber.e("Fail to restore session")
+            }
+    }
+
     @Suppress("UNCHECKED_CAST")
     fun restore(state: SavedStateMap?) {
         Timber.d("Restore state")
@@ -67,14 +93,7 @@ class MatrixClientsHolder @Inject constructor(private val authenticationService:
         // 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.d("Restore matrix session: $sessionId")
-                authenticationService.restoreSession(sessionId)
-                    .onSuccess { matrixClient ->
-                        add(matrixClient)
-                    }
-                    .onFailure {
-                        Timber.e("Fail to restore session")
-                    }
+                restore(sessionId)
             }
         }
     }

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/notifications/DefaultNotificationDrawerManager.kt

@@ -257,11 +257,11 @@ class DefaultNotificationDrawerManager @Inject constructor(
             val currentUser = tryOrNull(
                 onError = { Timber.e(it, "Unable to retrieve info for user ${sessionId.value}") },
                 operation = {
-                    val client = matrixClientsHolder.getOrNull(sessionId)
+                    val client = matrixClientsHolder.requireSession(sessionId)
 
                     // myUserDisplayName cannot be empty else NotificationCompat.MessagingStyle() will crash
-                    val myUserDisplayName = client?.loadUserDisplayName()?.getOrNull() ?: sessionId.value
-                    val userAvatarUrl = client?.loadUserAvatarURLString()?.getOrNull()
+                    val myUserDisplayName = client.loadUserDisplayName().getOrNull() ?: sessionId.value
+                    val userAvatarUrl = client.loadUserAvatarURLString().getOrNull()
                     MatrixUser(
                         userId = sessionId,
                         displayName = myUserDisplayName,

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/notifications/NotifiableEventResolver.kt

@@ -68,7 +68,7 @@ class NotifiableEventResolver @Inject constructor(
 
     suspend fun resolveEvent(sessionId: SessionId, roomId: RoomId, eventId: EventId): NotifiableEvent? {
         // Restore session
-        val client = matrixClientsHolder.getOrNull(sessionId) ?: return null
+        val client = matrixClientsHolder.requireSession(sessionId)
         val notificationService = client.notificationService()
         val notificationData = notificationService.getNotification(
                 userId = sessionId,