Fix media viewer flickering (#6800)

* Fix media viewer flickering

This was caused by the data being loaded triggering an index update that got out of sync with the displayed items, and for an instant, the pager index pointed to the wrong data until it was refreshed

* Reuse the same 'displayer' function for both forwards and backwards pagination

* Make `dataFlow` a property so we don't create a new instance every time we access it

* Remove `pageDataComparator` as it prevented new items from being loaded when a pagination returned no valid items to display but has more items to load

Make sure we modify the current index when loading new data only if it was pointing to the input event id.

* Fix `MediaViewerDataSource` overriding the provided timestamp for `Loading` items

Test emitting different loading items from the data source results in the state displaying the different items, so they will trigger a new pagination attempt

* Add regression test to check loading -> error -> loading states will still trigger 2 separate `LoadMore` events

libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/datasource/MediaGalleryDataSource.kt

@@ -32,6 +32,7 @@ import timber.log.Timber
 import java.util.concurrent.atomic.AtomicBoolean
 
 interface MediaGalleryDataSource {
+    val isReady: Boolean
     fun start(coroutineScope: CoroutineScope)
     fun groupedMediaItemsFlow(): Flow<AsyncData<GroupedMediaItems>>
     fun getLastData(): AsyncData<GroupedMediaItems>
@@ -49,7 +50,7 @@ class TimelineMediaGalleryDataSource(
 ) : MediaGalleryDataSource {
     private var timeline: Timeline? = null
 
-    private val groupedMediaItemsFlow = MutableSharedFlow<AsyncData<GroupedMediaItems>>(replay = 1)
+    private val groupedMediaItemsFlow = MutableSharedFlow<AsyncData<GroupedMediaItems>>(replay = 1, extraBufferCapacity = 10)
 
     override fun groupedMediaItemsFlow(): Flow<AsyncData<GroupedMediaItems>> = groupedMediaItemsFlow
 
@@ -59,9 +60,12 @@ class TimelineMediaGalleryDataSource(
 
     private val isStarted = AtomicBoolean(false)
 
+    override val isReady: Boolean get() = isStarted.get() && timeline != null
+
     @OptIn(ExperimentalCoroutinesApi::class)
     override fun start(coroutineScope: CoroutineScope) {
         if (!isStarted.compareAndSet(false, true)) {
+            Timber.w("MediaGalleryDataSource for room ${room.roomId} is already started, ignoring subsequent start call")
             return
         }
         flow {
@@ -73,10 +77,12 @@ class TimelineMediaGalleryDataSource(
             }
             mediaTimeline.getTimeline().fold(
                 {
+                    Timber.d("Timeline media data source flow started for room ${room.roomId}")
                     timeline = it
                     emit(it)
                 },
                 {
+                    Timber.e(it, "Failed to get media timeline for room ${room.roomId}")
                     groupedMediaItemsFlow.emit(AsyncData.Failure(it))
                 },
             )
@@ -107,13 +113,13 @@ class TimelineMediaGalleryDataSource(
     }
 
     override suspend fun loadMore(direction: Timeline.PaginationDirection) {
-        timeline?.paginate(direction)
+        timeline?.paginate(direction) ?: Timber.w("Timeline is not ready yet, cannot load more media items for room ${room.roomId}")
     }
 
     override suspend fun deleteItem(eventId: EventId) {
         timeline?.redactEvent(
             eventOrTransactionId = eventId.toEventOrTransactionId(),
             reason = null,
-        )
+        ) ?: Timber.w("Timeline is not ready yet, cannot delete media item with eventId $eventId for room ${room.roomId}")
     }
 }

libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/gallery/MediaGalleryPresenter.kt

@@ -94,7 +94,9 @@ class MediaGalleryPresenter(
                     mode = event.mode
                 }
                 is MediaGalleryEvent.LoadMore -> coroutineScope.launch {
-                    mediaGalleryDataSource.loadMore(event.direction)
+                    if (mediaGalleryDataSource.isReady) {
+                        mediaGalleryDataSource.loadMore(event.direction)
+                    }
                 }
                 is MediaGalleryEvent.Delete -> coroutineScope.launch {
                     mediaGalleryDataSource.deleteItem(event.eventId)

libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSource.kt

@@ -11,12 +11,13 @@ package io.element.android.libraries.mediaviewer.impl.viewer
 import androidx.annotation.VisibleForTesting
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.MutableState
+import androidx.compose.runtime.ProduceStateScope
 import androidx.compose.runtime.State
-import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.remember
+import androidx.compose.runtime.produceState
 import io.element.android.libraries.architecture.AsyncData
 import io.element.android.libraries.core.extensions.mapCatchingExceptions
+import io.element.android.libraries.matrix.api.core.EventId
 import io.element.android.libraries.matrix.api.media.MatrixMediaLoader
 import io.element.android.libraries.matrix.api.media.MediaFile
 import io.element.android.libraries.matrix.api.media.MediaSource
@@ -37,14 +38,17 @@ import kotlinx.collections.immutable.persistentListOf
 import kotlinx.collections.immutable.toImmutableList
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.withContext
 import timber.log.Timber
 import java.util.concurrent.ConcurrentHashMap
 
 class MediaViewerDataSource(
     mode: MediaViewerMode,
+    coroutineScope: CoroutineScope,
     private val dispatcher: CoroutineDispatcher,
     private val galleryDataSource: MediaGalleryDataSource,
     private val mediaLoader: MatrixMediaLoader,
@@ -76,40 +80,58 @@ class MediaViewerDataSource(
         localMediaStates.clear()
     }
 
+    /**
+     * Helper function to translate the [dataFlow] result to a Compose [State] that can be observed in the UI.
+     */
     @Composable
-    fun collectAsState(): State<ImmutableList<MediaViewerPageData>> {
-        return remember { dataFlow() }.collectAsState(initialData())
+    fun produceState(
+        producer: suspend ProduceStateScope<ImmutableList<MediaViewerPageData>>.(StateFlow<ImmutableList<MediaViewerPageData>>) -> Unit
+    ): State<ImmutableList<MediaViewerPageData>> {
+        return produceState(initialValue = initialData()) {
+            producer(dataFlow)
+        }
+    }
+
+    /**
+     * Find the index of the page corresponding to the given eventId, or null if not found.
+     */
+    fun findEventIndex(eventId: EventId?): Int? {
+        if (eventId == null) return null
+        return dataFlow.value.indexOfFirst { (it as? MediaViewerPageData.MediaViewerData)?.eventId == eventId }.takeIf { it >= 0 }
     }
 
     @VisibleForTesting
-    internal fun dataFlow(): Flow<ImmutableList<MediaViewerPageData>> {
-        return galleryDataSource.groupedMediaItemsFlow()
-            .map { groupedItems ->
-                when (groupedItems) {
-                    AsyncData.Uninitialized,
-                    is AsyncData.Loading -> {
-                        persistentListOf(
-                            MediaViewerPageData.Loading(
-                                direction = Timeline.PaginationDirection.BACKWARDS,
-                                timestamp = systemClock.epochMillis(),
-                                pagerKey = Long.MIN_VALUE,
-                            )
+    internal val dataFlow: StateFlow<ImmutableList<MediaViewerPageData>> = galleryDataSource.groupedMediaItemsFlow()
+        .map { groupedItems ->
+            when (groupedItems) {
+                AsyncData.Uninitialized,
+                is AsyncData.Loading -> {
+                    persistentListOf(
+                        MediaViewerPageData.Loading(
+                            direction = Timeline.PaginationDirection.BACKWARDS,
+                            timestamp = systemClock.epochMillis(),
+                            pagerKey = Long.MIN_VALUE,
                         )
-                    }
-                    is AsyncData.Failure -> {
-                        persistentListOf(
-                            MediaViewerPageData.Failure(groupedItems.error),
-                        )
-                    }
-                    is AsyncData.Success -> {
-                        withContext(dispatcher) {
-                            val mediaItems = groupedItems.data.getItems(galleryMode)
-                            buildMediaViewerPageList(mediaItems)
-                        }
+                    )
+                }
+                is AsyncData.Failure -> {
+                    persistentListOf(
+                        MediaViewerPageData.Failure(groupedItems.error),
+                    )
+                }
+                is AsyncData.Success -> {
+                    withContext(dispatcher) {
+                        val mediaItems = groupedItems.data.getItems(galleryMode)
+                        buildMediaViewerPageList(mediaItems)
                     }
                 }
             }
-    }
+        }
+        .stateIn(
+            scope = CoroutineScope(coroutineScope.coroutineContext + dispatcher),
+            started = SharingStarted.Lazily,
+            initialValue = initialData(),
+        )
 
     private fun initialData(): ImmutableList<MediaViewerPageData> {
         val initialMediaItems =
@@ -148,7 +170,7 @@ class MediaViewerDataSource(
                 is MediaItem.LoadingIndicator -> add(
                     MediaViewerPageData.Loading(
                         direction = mediaItem.direction,
-                        timestamp = systemClock.epochMillis(),
+                        timestamp = mediaItem.timestamp,
                         pagerKey = pagerKeysHandler.getKey(mediaItem),
                     )
                 )
@@ -161,7 +183,9 @@ class MediaViewerDataSource(
     }
 
     suspend fun loadMore(direction: Timeline.PaginationDirection) {
-        galleryDataSource.loadMore(direction)
+        if (galleryDataSource.isReady) {
+            galleryDataSource.loadMore(direction)
+        }
     }
 
     suspend fun loadMedia(data: MediaViewerPageData.MediaViewerData) {

libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerNode.kt

@@ -13,6 +13,7 @@ import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
+import androidx.lifecycle.lifecycleScope
 import com.bumble.appyx.core.modality.BuildContext
 import com.bumble.appyx.core.node.Node
 import com.bumble.appyx.core.plugin.Plugin
@@ -118,6 +119,7 @@ class MediaViewerNode(
         navigator = this,
         dataSource = MediaViewerDataSource(
             mode = inputs.mode,
+            coroutineScope = lifecycleScope,
             dispatcher = coroutineDispatchers.computation,
             galleryDataSource = mediaGallerySource,
             mediaLoader = mediaLoader,

libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerPresenter.kt

@@ -14,14 +14,12 @@ import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.IntState
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.State
-import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableIntStateOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.setValue
-import androidx.compose.runtime.snapshotFlow
 import dev.zacsweers.metro.Assisted
 import dev.zacsweers.metro.AssistedFactory
 import dev.zacsweers.metro.AssistedInject
@@ -45,10 +43,7 @@ import io.element.android.libraries.mediaviewer.impl.model.mediaPermissions
 import io.element.android.libraries.ui.strings.CommonStrings
 import kotlinx.collections.immutable.ImmutableList
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.filter
-import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.launch
 import io.element.android.libraries.androidutils.R as UtilsR
 
@@ -75,12 +70,27 @@ class MediaViewerPresenter(
     @Composable
     override fun present(): MediaViewerState {
         val coroutineScope = rememberCoroutineScope()
-        val data = dataSource.collectAsState()
-        val currentIndex = remember { mutableIntStateOf(searchIndex(data.value, inputs.eventId)) }
+        val currentIndex = remember { mutableIntStateOf(dataSource.findEventIndex(inputs.eventId) ?: 0) }
+        val data = dataSource.produceState { flow ->
+            flow.collectLatest { new ->
+                val existingItem = value.getOrNull(currentIndex.intValue)
+                val newItem = new.getOrNull(currentIndex.intValue)
+                if (existingItem is MediaViewerPageData.MediaViewerData && existingItem.eventId == inputs.eventId && newItem != existingItem) {
+                    currentIndex.intValue = dataSource.findEventIndex(inputs.eventId) ?: 0
+                } else if (currentIndex.intValue > 0 && value.firstOrNull() is MediaViewerPageData.Loading &&
+                    new.firstOrNull() !is MediaViewerPageData.Loading) {
+                    // Restore index based on the eventId after the initial items have been loaded
+                    currentIndex.intValue = dataSource.findEventIndex(inputs.eventId) ?: 0
+                }
+                value = new
+            }
+        }
+
         val snackbarMessage by snackbarDispatcher.collectSnackbarMessageAsState()
 
-        NoMoreItemsBackwardSnackBarDisplayer(currentIndex, data)
-        NoMoreItemsForwardSnackBarDisplayer(currentIndex, data)
+        // Add both forward and backward pagination state checks to display a snackbar when there is no more items to load in either direction
+        NoMoreItemsSnackBarDisplayer(currentIndex, data, Timeline.PaginationDirection.FORWARDS)
+        NoMoreItemsSnackBarDisplayer(currentIndex, data, Timeline.PaginationDirection.BACKWARDS)
 
         val permissions by room.permissionsAsState(MediaPermissions.DEFAULT) { perms ->
             perms.mediaPermissions()
@@ -176,52 +186,39 @@ class MediaViewerPresenter(
     }
 
     @Composable
-    private fun NoMoreItemsBackwardSnackBarDisplayer(
+    private fun NoMoreItemsSnackBarDisplayer(
         currentIndex: IntState,
         data: State<ImmutableList<MediaViewerPageData>>,
+        direction: Timeline.PaginationDirection,
     ) {
-        // With newest-first ordering, backward loading indicator is at the last index
-        val isRenderingLoadingBackward by remember {
-            derivedStateOf {
-                currentIndex.intValue == data.value.lastIndex &&
-                    data.value.size > 1 &&
-                    data.value.lastOrNull() is MediaViewerPageData.Loading
+        var previousIndex by remember { mutableIntStateOf(currentIndex.intValue) }
+        var previousDataSize by remember { mutableIntStateOf(data.value.size) }
+        var wasLoading: Boolean? by remember { mutableStateOf(null) }
+        LaunchedEffect(currentIndex.intValue, data.value) {
+            fun isLoading(index: Int, data: List<MediaViewerPageData>, direction: Timeline.PaginationDirection): Boolean {
+                return when (direction) {
+                    Timeline.PaginationDirection.BACKWARDS -> index == data.lastIndex && data.lastOrNull() is MediaViewerPageData.Loading
+                    Timeline.PaginationDirection.FORWARDS -> index == 0 && data.firstOrNull() is MediaViewerPageData.Loading
+                }
             }
-        }
-        if (isRenderingLoadingBackward) {
-            LaunchedEffect(Unit) {
-                // Observe the loading data vanishing
-                snapshotFlow { data.value.lastOrNull() is MediaViewerPageData.Loading }
-                    .distinctUntilChanged()
-                    .filter { !it }
-                    .onEach { showNoMoreItemsSnackbar() }
-                    .launchIn(this)
+            // Reset the effect when the user navigate to another item so we only take into account index changes caused by data changes
+            if (previousIndex != currentIndex.intValue) {
+                wasLoading = null
+                previousIndex = currentIndex.intValue
+            }
+            // If we were navigating backwards and the data size grew, we can discard the previous value: it means we received new items
+            if (direction == Timeline.PaginationDirection.BACKWARDS && previousDataSize < data.value.size) {
+                wasLoading = null
             }
-        }
-    }
 
-    @Composable
-    private fun NoMoreItemsForwardSnackBarDisplayer(
-        currentIndex: IntState,
-        data: State<ImmutableList<MediaViewerPageData>>,
-    ) {
-        // With newest-first ordering, forward loading indicator is at the first index
-        val isRenderingLoadingForward by remember {
-            derivedStateOf {
-                currentIndex.intValue == 0 &&
-                    data.value.size > 1 &&
-                    data.value.firstOrNull() is MediaViewerPageData.Loading
-            }
-        }
-        if (isRenderingLoadingForward) {
-            LaunchedEffect(Unit) {
-                // Observe the loading data vanishing
-                snapshotFlow { data.value.firstOrNull() is MediaViewerPageData.Loading }
-                    .distinctUntilChanged()
-                    .filter { !it }
-                    .onEach { showNoMoreItemsSnackbar() }
-                    .launchIn(this)
+            val isLoading = isLoading(currentIndex.intValue, data.value, direction)
+
+            if (wasLoading == true && !isLoading) {
+                showNoMoreItemsSnackbar()
             }
+
+            previousDataSize = data.value.size
+            wasLoading = isLoading
         }
     }
 
@@ -293,13 +290,4 @@ class MediaViewerPresenter(
             CommonStrings.error_unknown
         }
     }
-
-    private fun searchIndex(data: List<MediaViewerPageData>, eventId: EventId?): Int {
-        if (eventId == null) {
-            return 0
-        }
-        return data.indexOfFirst {
-            (it as? MediaViewerPageData.MediaViewerData)?.eventId == eventId
-        }.coerceAtLeast(0)
-    }
 }

libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerState.kt

@@ -8,6 +8,7 @@
 
 package io.element.android.libraries.mediaviewer.impl.viewer
 
+import androidx.compose.runtime.Immutable
 import androidx.compose.runtime.State
 import io.element.android.libraries.architecture.AsyncData
 import io.element.android.libraries.designsystem.utils.snackbar.SnackbarMessage
@@ -29,6 +30,7 @@ data class MediaViewerState(
     val eventSink: (MediaViewerEvent) -> Unit,
 )
 
+@Immutable
 sealed interface MediaViewerPageData {
     val pagerKey: Long
 

libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerView.kt

@@ -44,7 +44,6 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.runtime.setValue
-import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.alpha
@@ -179,9 +178,11 @@ fun MediaViewerView(
         val pagerState = rememberPagerState(state.currentIndex, 0f) {
             state.listData.size
         }
-        LaunchedEffect(pagerState) {
-            snapshotFlow { pagerState.currentPage }.collect { page ->
-                state.eventSink(MediaViewerEvent.OnNavigateTo(page))
+
+        LaunchedEffect(pagerState.targetPage, state.currentIndex) {
+            // Only emit an index navigation change when it's triggered by the user scrolling
+            if (pagerState.targetPage != state.currentIndex && pagerState.isScrollInProgress) {
+                state.eventSink(MediaViewerEvent.OnNavigateTo(pagerState.targetPage))
             }
         }
         HorizontalPager(

libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/SingleMediaGalleryDataSource.kt

@@ -30,6 +30,8 @@ class SingleMediaGalleryDataSource(
     override fun groupedMediaItemsFlow() = flowOf(AsyncData.Success(data))
     override fun getLastData(): AsyncData<GroupedMediaItems> = AsyncData.Success(data)
 
+    override val isReady: Boolean = true
+
     override suspend fun loadMore(direction: Timeline.PaginationDirection) = Unit
     override suspend fun deleteItem(eventId: EventId) = Unit