feat: handle cell urls expiration (WPB-21328) (#4417)

Author: sbakhtiarov

app/src/main/kotlin/com/wire/android/ui/calling/common/SharedCallingViewModel.kt

@@ -338,7 +338,7 @@ class SharedCallingViewModel @AssistedInject constructor(
     private fun recentInCallReactionMap(): MutableMap<UserId, String> =
         ExpiringMap<UserId, String>(
             scope = viewModelScope,
-            expiration = InCallReactions.recentReactionShowDurationMs,
+            expirationMs = InCallReactions.recentReactionShowDurationMs,
             delegate = mutableStateMapOf<UserId, String>()
         )
 

app/src/main/kotlin/com/wire/android/ui/common/multipart/MultipartAttachmentUi.kt

@@ -31,6 +31,7 @@ data class MultipartAttachmentUi(
     val localPath: String?,
     val contentHash: String? = null,
     val contentUrl: String? = null,
+    val contentUrlExpiresAt: Long? = null,
     val previewUrl: String? = null,
     val mimeType: String,
     val assetType: AttachmentFileType,
@@ -55,6 +56,7 @@ fun CellAssetContent.toUiModel(progress: Float?) = MultipartAttachmentUi(
     fileName = this.assetPath?.substringAfterLast("/"),
     localPath = this.localPath,
     contentUrl = this.contentUrl,
+    contentUrlExpiresAt = this.contentUrlExpiresAt,
     previewUrl = this.previewUrl,
     mimeType = this.mimeType,
     assetType = AttachmentFileType.fromMimeType(mimeType),

app/src/main/kotlin/com/wire/android/ui/home/conversations/model/messagetypes/multipart/MultipartAttachmentsViewModel.kt

@@ -21,13 +21,14 @@ import androidx.compose.runtime.mutableStateMapOf
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.wire.android.appLogger
+import com.wire.android.feature.cells.domain.model.AttachmentFileType
 import com.wire.android.feature.cells.domain.model.AttachmentFileType.IMAGE
 import com.wire.android.feature.cells.domain.model.AttachmentFileType.PDF
 import com.wire.android.feature.cells.domain.model.AttachmentFileType.VIDEO
-import com.wire.android.feature.cells.domain.model.AttachmentFileType
 import com.wire.android.ui.common.multipart.AssetSource
 import com.wire.android.ui.common.multipart.MultipartAttachmentUi
 import com.wire.android.ui.common.multipart.toUiModel
+import com.wire.android.util.ExpiringMap
 import com.wire.android.util.FileManager
 import com.wire.kalium.cells.domain.usecase.DownloadCellFileUseCase
 import com.wire.kalium.cells.domain.usecase.RefreshCellAssetStateUseCase
@@ -41,6 +42,7 @@ import kotlinx.collections.immutable.toImmutableList
 import kotlinx.coroutines.launch
 import okio.Path.Companion.toPath
 import javax.inject.Inject
+import kotlin.time.Duration.Companion.hours
 
 @HiltViewModel
 class MultipartAttachmentsViewModel @Inject constructor(
@@ -50,7 +52,18 @@ class MultipartAttachmentsViewModel @Inject constructor(
     private val kaliumFileSystem: KaliumFileSystem,
 ) : ViewModel() {
 
-    private val refreshed = mutableListOf<String>()
+    private companion object {
+        val DEFAULT_CONTENT_URL_EXPIRY_MS = 1.hours.inWholeMilliseconds
+    }
+
+    private val refreshed = ExpiringMap<String, Unit>(
+        scope = viewModelScope,
+        expirationMs = DEFAULT_CONTENT_URL_EXPIRY_MS,
+        delegate = mutableMapOf(),
+        onEntryExpired = { key, _ ->
+            viewModelScope.launch { refreshAsset(key) }
+        }
+    )
 
     private val uploadProgress = mutableStateMapOf<String, Float>()
 
@@ -105,12 +118,17 @@ class MultipartAttachmentsViewModel @Inject constructor(
     }
 
     fun refreshAssetState(attachment: MultipartAttachmentUi) {
-        if (refreshed.contains(attachment.uuid).not()) {
-            refreshed.add(attachment.uuid)
-            if (attachment.source == AssetSource.CELL) {
-                viewModelScope.launch { refreshAsset(attachment.uuid) }
-            }
+
+        if (attachment.source != AssetSource.CELL) return
+        if (refreshed.contains(attachment.uuid)) return
+
+        if (attachment.contentUrlExpiresAt != null) {
+            refreshed.putWithExpireAt(attachment.uuid, Unit, attachment.contentUrlExpiresAt)
+        } else {
+            refreshed.put(attachment.uuid, Unit)
         }
+
+        viewModelScope.launch { refreshAsset(attachment.uuid) }
     }
 
     private fun openLocalFile(attachment: MultipartAttachmentUi) {

app/src/main/kotlin/com/wire/android/util/ExpiringMap.kt

@@ -28,17 +28,25 @@ import java.util.concurrent.ConcurrentHashMap
  */
 class ExpiringMap<K, V>(
     private val scope: CoroutineScope,
-    private val expiration: Long,
+    private val expirationMs: Long,
     private val delegate: MutableMap<K, V>,
+    private val onEntryExpired: ((key: K, value: V?) -> Unit)? = null,
     private val currentTime: () -> Long = { System.currentTimeMillis() },
 ) : MutableMap<K, V> by delegate {
 
     private val timestamps: MutableMap<K, Long> = ConcurrentHashMap<K, Long>()
     private var cleanupJob: Job? = null
 
+    fun putWithExpireAt(key: K, value: V, expireAt: Long): V? {
+        return delegate.put(key, value).also {
+            timestamps.put(key, expireAt)
+            scheduleCleanup()
+        }
+    }
+
     override fun put(key: K, value: V): V? {
         return delegate.put(key, value).also {
-            timestamps.put(key, currentTime() + expiration)
+            timestamps.put(key, currentTime() + expirationMs)
             scheduleCleanup()
         }
     }
@@ -65,7 +73,8 @@ class ExpiringMap<K, V>(
         val now = currentTime()
         timestamps.entries.onEach { (key, expiration) ->
             if (expiration <= now) {
-                delegate.remove(key)
+                val value = delegate.remove(key)
+                onEntryExpired?.invoke(key, value)
             }
         }
         timestamps.entries.removeAll { it.value <= now }

app/src/test/kotlin/com/wire/android/util/ExpiringMapTest.kt

@@ -52,7 +52,7 @@ class ExpiringMapTest {
     }
 
     @Test
-    fun `check item can not be obtained before expiration`() = runTest {
+    fun `check item can be obtained before expiration`() = runTest {
         // given
         val map = withTestExpiringMap()
 
@@ -107,9 +107,22 @@ class ExpiringMapTest {
         assertEquals(null, map["testKey"])
     }
 
-    private fun TestScope.withTestExpiringMap(): MutableMap<String, String> = ExpiringMap<String, String>(
+    @Test
+    fun `check item can be put with custom expiration`() = runTest {
+        // given
+        val map = withTestExpiringMap()
+
+        // when
+        map.putWithExpireAt("testKey", "testValue", currentTime + 500)
+        advanceTimeBy(301) // advance past default expiration
+
+        // then
+        assertEquals("testValue", map["testKey"])
+    }
+
+    private fun TestScope.withTestExpiringMap(): ExpiringMap<String, String> = ExpiringMap<String, String>(
         scope = this.backgroundScope,
-        expiration = 300,
+        expirationMs = 300,
         delegate = mutableMapOf(),
         currentTime = { currentTime }
     )