improve the profile cache

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,27 +18,41 @@ 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 {
-            nostr.waitUntilInitialized()
-            loadCacheMetadata()
+            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()
+                getCachedMetadata()
+            } catch (e: Exception) {
+                showError("Failed to load initial cache: ${e.message}")
+            }
         }
     }
 
-    private suspend fun runObserver() = coroutineScope {
-        launch {
-            nostr.profiles.metadataUpdates.collect { (pubkey, metadata) ->
-                updateMetadata(pubkey, Profile(pubkey, metadata))
-            }
+    private suspend fun runObserver() {
+        nostr.profiles.metadataUpdates.collect { (pubkey, metadata) ->
+            updateMetadata(pubkey, Profile(pubkey, metadata))
         }
     }
 
@@ -57,37 +70,69 @@ class ProfileCache(
                 batch.add(nextKey)
             }
 
-            nostr.profiles.fetchMetadataBatch(batch.toList())
-        }
-    }
-
-    private suspend fun loadCacheMetadata() {
-        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)
+            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 fun requestMetadata(pubkey: PublicKey) {
-        if (seenPublicKeys.add(pubkey)) {
-            metadataRequestChannel.trySend(pubkey)
+    private suspend fun getCachedMetadata() {
+        val cache = nostr.profiles.getAllCacheMetadata()
+        cache.forEach { (pubkey, metadata) ->
+            updateMetadata(pubkey, Profile(pubkey, metadata))
         }
     }
 
-    private suspend fun updateMetadata(pubkey: PublicKey, profile: Profile) {
-        profilesMutex.withLock {
-            profiles.getOrPut(pubkey) { MutableStateFlow(null) }.value = profile
+    private fun requestMetadata(pubkey: PublicKey) {
+        var added = false
+        seenPublicKeys.update { current ->
+            if (current.contains(pubkey)) {
+                added = false
+                current
+            } else {
+                added = true
+                current + pubkey
+            }
         }
+        if (added) metadataRequestChannel.trySend(pubkey)
+    }
+
+    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()
     }
 }