improve the profile cache

composeApp/src/androidMain/kotlin/su/reya/coop/screens/HomeScreen.kt

@@ -98,9 +98,9 @@ import su.reya.coop.RoomKind
 import su.reya.coop.RoomUiState
 import su.reya.coop.Screen
 import su.reya.coop.ago
+import su.reya.coop.flow
 import su.reya.coop.shared.Avatar
 import su.reya.coop.shared.getExpressiveFontFamily
-import su.reya.coop.uiStateFlow
 import su.reya.coop.viewmodel.AccountViewModel
 import su.reya.coop.viewmodel.ChatViewModel
 
@@ -620,9 +620,9 @@ fun NewRequests(requests: List<Room>) {
     val firstRoom = requests.getOrNull(0)
     val secondRoom = requests.getOrNull(1)
 
-    val firstRoomState by (firstRoom as Room).uiStateFlow(profileCache)
+    val firstRoomState by (firstRoom as Room).flow(profileCache)
         .collectAsStateWithLifecycle(RoomUiState())
-    val secondRoomState by (secondRoom ?: firstRoom).uiStateFlow(profileCache)
+    val secondRoomState by (secondRoom ?: firstRoom).flow(profileCache)
         .collectAsStateWithLifecycle(RoomUiState())
 
     val supportingText = when {
@@ -695,8 +695,7 @@ fun NewRequests(requests: List<Room>) {
 @Composable
 fun ChatRoom(room: Room, onClick: () -> Unit) {
     val profileCache = LocalProfileCache.current
-    val roomState by room.uiStateFlow(profileCache)
-        .collectAsStateWithLifecycle(RoomUiState())
+    val roomState by room.flow(profileCache).collectAsStateWithLifecycle(RoomUiState())
 
     ListItem(
         modifier = Modifier.clickable(onClick = onClick),

shared/src/commonMain/kotlin/su/reya/coop/viewmodel/AccountViewModel.kt

@@ -13,7 +13,6 @@ import su.reya.coop.repository.AccountState
 class AccountViewModel(
     private val repository: AccountRepository,
 ) : ViewModel(), ErrorHost by repository {
-
     val state: StateFlow<AccountState> = repository.state
     val isUpdatingProfile: StateFlow<Boolean> = repository.isUpdatingProfile
     val currentUserProfile: StateFlow<Profile?> = repository.currentUserProfile

shared/src/commonMain/kotlin/su/reya/coop/viewmodel/ProfileCache.kt

@@ -1,16 +1,15 @@
 package su.reya.coop.viewmodel
 
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.sync.Mutex
-import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withTimeoutOrNull
 import rust.nostr.sdk.PublicKey
 import su.reya.coop.Profile
@@ -19,29 +18,43 @@ import kotlin.time.Duration.Companion.milliseconds
 
 class ProfileCache(
     private val nostr: Nostr,
+    private val defaultDispatcher: CoroutineDispatcher = Dispatchers.Default,
 ) : ErrorHost by createErrorHost() {
-    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
-    private val profilesMutex = Mutex()
-    private val profiles = mutableMapOf<PublicKey, MutableStateFlow<Profile?>>()
+    private val scope = CoroutineScope(SupervisorJob() + defaultDispatcher)
+    private val profiles = MutableStateFlow<Map<PublicKey, MutableStateFlow<Profile?>>>(emptyMap())
+    private val seenPublicKeys = MutableStateFlow<Set<PublicKey>>(emptySet())
     private val metadataRequestChannel = Channel<PublicKey>(Channel.UNLIMITED)
-    private val seenPublicKeys = mutableSetOf<PublicKey>()
 
     init {
-        scope.launch { runObserver() }
-        scope.launch { runMetadataBatching() }
         scope.launch {
+            try {
+                runObserver()
+            } catch (e: Exception) {
+                showError("Metadata observer failed: ${e.message}")
+            }
+        }
+        scope.launch {
+            try {
+                runMetadataBatching()
+            } catch (e: Exception) {
+                showError("Metadata batching failed: ${e.message}")
+            }
+        }
+        scope.launch {
+            try {
                 nostr.waitUntilInitialized()
-            loadCacheMetadata()
+                getCachedMetadata()
+            } catch (e: Exception) {
+                showError("Failed to load initial cache: ${e.message}")
+            }
         }
     }
 
-    private suspend fun runObserver() = coroutineScope {
-        launch {
+    private suspend fun runObserver() {
         nostr.profiles.metadataUpdates.collect { (pubkey, metadata) ->
             updateMetadata(pubkey, Profile(pubkey, metadata))
         }
     }
-    }
 
     private suspend fun runMetadataBatching() {
         nostr.waitUntilInitialized()
@@ -57,37 +70,69 @@ class ProfileCache(
                 batch.add(nextKey)
             }
 
+            try {
                 nostr.profiles.fetchMetadataBatch(batch.toList())
+            } catch (e: Exception) {
+                // Allow these keys to be requested again since the fetch failed
+                seenPublicKeys.update { it - batch }
+                println("Failed to fetch metadata batch: ${e.message}")
+            }
         }
     }
 
-    private suspend fun loadCacheMetadata() {
+    private suspend fun getCachedMetadata() {
         val cache = nostr.profiles.getAllCacheMetadata()
-
-        profilesMutex.withLock {
         cache.forEach { (pubkey, metadata) ->
-                val profile = Profile(pubkey, metadata)
-                profiles.getOrPut(pubkey) { MutableStateFlow(null) }.value = profile
-                seenPublicKeys.add(pubkey)
-            }
+            updateMetadata(pubkey, Profile(pubkey, metadata))
         }
     }
 
     private fun requestMetadata(pubkey: PublicKey) {
-        if (seenPublicKeys.add(pubkey)) {
-            metadataRequestChannel.trySend(pubkey)
+        var added = false
+        seenPublicKeys.update { current ->
+            if (current.contains(pubkey)) {
+                added = false
+                current
+            } else {
+                added = true
+                current + pubkey
             }
         }
+        if (added) metadataRequestChannel.trySend(pubkey)
+    }
 
-    private suspend fun updateMetadata(pubkey: PublicKey, profile: Profile) {
-        profilesMutex.withLock {
-            profiles.getOrPut(pubkey) { MutableStateFlow(null) }.value = profile
+    private fun updateMetadata(pubkey: PublicKey, profile: Profile) {
+        profiles.update { current ->
+            val flow = current[pubkey] ?: MutableStateFlow<Profile?>(null)
+            flow.value = profile
+            if (current.containsKey(pubkey)) current else current + (pubkey to flow)
         }
+        seenPublicKeys.update { it + pubkey }
     }
 
     fun getMetadata(pubkey: PublicKey): StateFlow<Profile?> {
-        val flow = profiles.getOrPut(pubkey) { MutableStateFlow(null) }
-        if (flow.value == null) requestMetadata(pubkey)
-        return flow.asStateFlow()
+        val currentMap = profiles.value
+        val existingFlow = currentMap[pubkey]
+
+        if (existingFlow != null) {
+            if (existingFlow.value == null) requestMetadata(pubkey)
+            return existingFlow.asStateFlow()
+        }
+
+        val newFlow = MutableStateFlow<Profile?>(null)
+        var resultFlow = newFlow
+
+        profiles.update { prev ->
+            if (prev.containsKey(pubkey)) {
+                resultFlow = prev[pubkey]!!
+                prev
+            } else {
+                prev + (pubkey to newFlow)
+            }
+        }
+
+        if (resultFlow.value == null) requestMetadata(pubkey)
+
+        return resultFlow.asStateFlow()
     }
 }