Unified preview image loader.

Replace ImagePreviewImageLoader and CachingPreviewImageLoader with the
new class, PreviewImageLoader.

Fix: 348665058
Fix: 343819590
Test: atest IntentResolver-tests-unit
Test: atest IntentResolver-tests-activity
Test: inject debugging logs for image requests, loadings, and
cancellations and verify common scenarios.
Flag: com.android.intentresolver.preview_image_loader

Change-Id: Ib54ea0d538cd2cbaef6041f09eeb2b9eb6ada4cf

5 files changed, 332 insertions, 94 deletions

diff --git a/java/src/com/android/intentresolver/contentpreview/CachingImagePreviewImageLoader.kt b/java/src/com/android/intentresolver/contentpreview/CachingImagePreviewImageLoader.kt
index f60f550e..847fcc82 100644
--- a/java/src/com/android/intentresolver/contentpreview/CachingImagePreviewImageLoader.kt
+++ b/java/src/com/android/intentresolver/contentpreview/CachingImagePreviewImageLoader.kt
@@ -87,7 +87,7 @@ constructor(
 
     private suspend fun loadUncachedImage(uri: Uri): Bitmap? =
         withContext(bgDispatcher) {
-            runCatching { semaphore.withPermit { thumbnailLoader.invoke(uri) } }
+            runCatching { semaphore.withPermit { thumbnailLoader.loadThumbnail(uri) } }
                 .onFailure {
                     ensureActive()
                     Log.d(TAG, "Failed to load preview for $uri", it)
diff --git a/java/src/com/android/intentresolver/contentpreview/ImageLoaderModule.kt b/java/src/com/android/intentresolver/contentpreview/ImageLoaderModule.kt
index 17d05099..27e817db 100644
--- a/java/src/com/android/intentresolver/contentpreview/ImageLoaderModule.kt
+++ b/java/src/com/android/intentresolver/contentpreview/ImageLoaderModule.kt
@@ -17,6 +17,7 @@
 package com.android.intentresolver.contentpreview
 
 import android.content.res.Resources
+import com.android.intentresolver.Flags
 import com.android.intentresolver.R
 import com.android.intentresolver.inject.ApplicationOwned
 import dagger.Binds
@@ -24,16 +25,26 @@ import dagger.Module
 import dagger.Provides
 import dagger.hilt.InstallIn
 import dagger.hilt.android.components.ViewModelComponent
+import javax.inject.Provider
 
 @Module
 @InstallIn(ViewModelComponent::class)
 interface ImageLoaderModule {
-    @Binds fun imageLoader(previewImageLoader: ImagePreviewImageLoader): ImageLoader
-
     @Binds fun thumbnailLoader(thumbnailLoader: ThumbnailLoaderImpl): ThumbnailLoader
 
     companion object {
         @Provides
+        fun imageLoader(
+            imagePreviewImageLoader: Provider<ImagePreviewImageLoader>,
+            previewImageLoader: Provider<PreviewImageLoader>
+        ): ImageLoader =
+            if (Flags.previewImageLoader()) {
+                previewImageLoader.get()
+            } else {
+                imagePreviewImageLoader.get()
+            }
+
+        @Provides
         @ThumbnailSize
         fun thumbnailSize(@ApplicationOwned resources: Resources): Int =
             resources.getDimensionPixelSize(R.dimen.chooser_preview_image_max_dimen)
diff --git a/java/src/com/android/intentresolver/contentpreview/PreviewImageLoader.kt b/java/src/com/android/intentresolver/contentpreview/PreviewImageLoader.kt
new file mode 100644
index 00000000..b10f7ef9
--- /dev/null
+++ b/java/src/com/android/intentresolver/contentpreview/PreviewImageLoader.kt
@@ -0,0 +1,197 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.intentresolver.contentpreview
+
+import android.graphics.Bitmap
+import android.net.Uri
+import android.util.Log
+import android.util.Size
+import androidx.collection.lruCache
+import com.android.intentresolver.inject.Background
+import com.android.intentresolver.inject.ViewModelOwned
+import javax.annotation.concurrent.GuardedBy
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.firstOrNull
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.sync.Semaphore
+import kotlinx.coroutines.sync.withPermit
+
+private const val TAG = "PayloadSelImageLoader"
+
+/**
+ * Implements preview image loading for the payload selection UI. Cancels preview loading for items
+ * that has been evicted from the cache at the expense of a possible request duplication (deemed
+ * unlikely).
+ */
+class PreviewImageLoader
+@Inject
+constructor(
+    @ViewModelOwned private val scope: CoroutineScope,
+    @PreviewCacheSize private val cacheSize: Int,
+    @ThumbnailSize private val defaultPreviewSize: Int,
+    private val thumbnailLoader: ThumbnailLoader,
+    @Background private val bgDispatcher: CoroutineDispatcher,
+    @PreviewMaxConcurrency maxSimultaneousRequests: Int = 4,
+) : ImageLoader {
+
+    private val contentResolverSemaphore = Semaphore(maxSimultaneousRequests)
+
+    private val lock = Any()
+    @GuardedBy("lock") private val runningRequests = hashMapOf<Uri, RequestRecord>()
+    @GuardedBy("lock")
+    private val cache =
+        lruCache<Uri, RequestRecord>(
+            maxSize = cacheSize,
+            onEntryRemoved = { _, _, oldRec, newRec ->
+                if (oldRec !== newRec) {
+                    onRecordEvictedFromCache(oldRec)
+                }
+            }
+        )
+
+    override suspend fun invoke(uri: Uri, size: Size, caching: Boolean): Bitmap? =
+        loadImageInternal(uri, size, caching)
+
+    override fun prePopulate(uriSizePairs: List<Pair<Uri, Size>>) {
+        uriSizePairs.asSequence().take(cacheSize).forEach { uri ->
+            scope.launch { loadImageInternal(uri.first, uri.second, caching = true) }
+        }
+    }
+
+    private suspend fun loadImageInternal(uri: Uri, size: Size, caching: Boolean): Bitmap? {
+        return withRequestRecord(uri, caching) { record ->
+            val newSize = sanitize(size)
+            val newMetric = newSize.metric
+            record
+                .also {
+                    // set the requested size to the max of the new and the previous value; input
+                    // will emit if the resulted value is greater than the old one
+                    it.input.update { oldSize ->
+                        if (oldSize == null || oldSize.metric < newSize.metric) newSize else oldSize
+                    }
+                }
+                .output
+                // filter out bitmaps of a lower resolution than that we're requesting
+                .filter { it is BitmapLoadingState.Loaded && newMetric <= it.size.metric }
+                .firstOrNull()
+                ?.let { (it as BitmapLoadingState.Loaded).bitmap }
+        }
+    }
+
+    private suspend fun withRequestRecord(
+        uri: Uri,
+        caching: Boolean,
+        block: suspend (RequestRecord) -> Bitmap?
+    ): Bitmap? {
+        val record = trackRecordRunning(uri, caching)
+        return try {
+            block(record)
+        } finally {
+            untrackRecordRunning(uri, record)
+        }
+    }
+
+    private fun trackRecordRunning(uri: Uri, caching: Boolean): RequestRecord =
+        synchronized(lock) {
+            runningRequests
+                .getOrPut(uri) { cache[uri] ?: createRecord(uri) }
+                .also { record ->
+                    record.clientCount++
+                    if (caching) {
+                        cache.put(uri, record)
+                    }
+                }
+        }
+
+    private fun untrackRecordRunning(uri: Uri, record: RequestRecord) {
+        synchronized(lock) {
+            record.clientCount--
+            if (record.clientCount <= 0) {
+                runningRequests.remove(uri)
+                val result = record.output.value
+                if (cache[uri] == null) {
+                    record.loadingJob.cancel()
+                } else if (result is BitmapLoadingState.Loaded && result.bitmap == null) {
+                    cache.remove(uri)
+                }
+            }
+        }
+    }
+
+    private fun onRecordEvictedFromCache(record: RequestRecord) {
+        synchronized(lock) {
+            if (record.clientCount <= 0) {
+                record.loadingJob.cancel()
+            }
+        }
+    }
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    private fun createRecord(uri: Uri): RequestRecord {
+        // use a StateFlow with sentinel values to avoid using SharedFlow that is deemed dangerous
+        val input = MutableStateFlow<Size?>(null)
+        val output = MutableStateFlow<BitmapLoadingState>(BitmapLoadingState.Loading)
+        val job =
+            scope.launch(bgDispatcher) {
+                // the image loading pipeline: input -- a desired image size, output -- a bitmap
+                input
+                    .filterNotNull()
+                    .mapLatest { size -> BitmapLoadingState.Loaded(size, loadBitmap(uri, size)) }
+                    .collect { output.tryEmit(it) }
+            }
+        return RequestRecord(input, output, job, clientCount = 0)
+    }
+
+    private suspend fun loadBitmap(uri: Uri, size: Size): Bitmap? =
+        contentResolverSemaphore.withPermit {
+            runCatching { thumbnailLoader.loadThumbnail(uri, size) }
+                .onFailure { Log.d(TAG, "failed to load $uri preview", it) }
+                .getOrNull()
+        }
+
+    private class RequestRecord(
+        /** The image loading pipeline input: desired preview size */
+        val input: MutableStateFlow<Size?>,
+        /** The image loading pipeline output */
+        val output: MutableStateFlow<BitmapLoadingState>,
+        /** The image loading pipeline job */
+        val loadingJob: Job,
+        @GuardedBy("lock") var clientCount: Int,
+    )
+
+    private sealed interface BitmapLoadingState {
+        data object Loading : BitmapLoadingState
+
+        data class Loaded(val size: Size, val bitmap: Bitmap?) : BitmapLoadingState
+    }
+
+    private fun sanitize(size: Size?): Size =
+        size?.takeIf { it.width > 0 && it.height > 0 }
+            ?: Size(defaultPreviewSize, defaultPreviewSize)
+}
+
+private val Size.metric
+    get() = maxOf(width, height)
diff --git a/java/src/com/android/intentresolver/contentpreview/ThumbnailLoader.kt b/java/src/com/android/intentresolver/contentpreview/ThumbnailLoader.kt
index 9f1d50da..e8afa480 100644
--- a/java/src/com/android/intentresolver/contentpreview/ThumbnailLoader.kt
+++ b/java/src/com/android/intentresolver/contentpreview/ThumbnailLoader.kt
@@ -20,10 +20,25 @@ import android.content.ContentResolver
 import android.graphics.Bitmap
 import android.net.Uri
 import android.util.Size
+import com.android.intentresolver.util.withCancellationSignal
 import javax.inject.Inject
 
 /** Interface for objects that can attempt load a [Bitmap] from a [Uri]. */
-interface ThumbnailLoader : suspend (Uri) -> Bitmap?
+interface ThumbnailLoader {
+    /**
+     * Loads a thumbnail for the given [uri].
+     *
+     * The size of the thumbnail is determined by the implementation.
+     */
+    suspend fun loadThumbnail(uri: Uri): Bitmap?
+
+    /**
+     * Loads a thumbnail for the given [uri] and [size].
+     *
+     * The [size] is the size of the thumbnail in pixels.
+     */
+    suspend fun loadThumbnail(uri: Uri, size: Size): Bitmap?
+}
 
 /** Default implementation of [ThumbnailLoader]. */
 class ThumbnailLoaderImpl
@@ -35,6 +50,11 @@ constructor(
 
     private val size = Size(thumbnailSize, thumbnailSize)
 
-    override suspend fun invoke(uri: Uri): Bitmap =
-        contentResolver.loadThumbnail(uri, size, /* signal = */ null)
+    override suspend fun loadThumbnail(uri: Uri): Bitmap =
+        contentResolver.loadThumbnail(uri, size, /* signal= */ null)
+
+    override suspend fun loadThumbnail(uri: Uri, size: Size): Bitmap =
+        withCancellationSignal { signal ->
+            contentResolver.loadThumbnail(uri, size, signal)
+        }
 }
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModel.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModel.kt
index f1e65f73..6f8be1ff 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModel.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModel.kt
@@ -16,10 +16,12 @@
 package com.android.intentresolver.contentpreview.payloadtoggle.ui.viewmodel
 
 import android.util.Size
+import com.android.intentresolver.Flags
 import com.android.intentresolver.contentpreview.CachingImagePreviewImageLoader
 import com.android.intentresolver.contentpreview.HeadlineGenerator
 import com.android.intentresolver.contentpreview.ImageLoader
 import com.android.intentresolver.contentpreview.MimeTypeClassifier
+import com.android.intentresolver.contentpreview.PreviewImageLoader
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.cursor.PayloadToggle
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.interactor.ChooserRequestInteractor
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.interactor.CustomActionsInteractor
@@ -30,11 +32,11 @@ import com.android.intentresolver.contentpreview.payloadtoggle.shared.ContentTyp
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewsModel
 import com.android.intentresolver.inject.ViewModelOwned
-import dagger.Binds
 import dagger.Module
 import dagger.Provides
 import dagger.hilt.InstallIn
 import dagger.hilt.android.components.ViewModelComponent
+import javax.inject.Provider
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
@@ -65,98 +67,106 @@ data class ShareouselViewModel(
 
 @Module
 @InstallIn(ViewModelComponent::class)
-interface ShareouselViewModelModule {
+object ShareouselViewModelModule {
 
-    @Binds @PayloadToggle fun imageLoader(imageLoader: CachingImagePreviewImageLoader): ImageLoader
+    @Provides
+    @PayloadToggle
+    fun imageLoader(
+        cachingImageLoader: Provider<CachingImagePreviewImageLoader>,
+        previewImageLoader: Provider<PreviewImageLoader>
+    ): ImageLoader =
+        if (Flags.previewImageLoader()) {
+            previewImageLoader.get()
+        } else {
+            cachingImageLoader.get()
+        }
 
-    companion object {
-        @Provides
-        fun create(
-            interactor: SelectablePreviewsInteractor,
-            @PayloadToggle imageLoader: ImageLoader,
-            actionsInteractor: CustomActionsInteractor,
-            headlineGenerator: HeadlineGenerator,
-            selectionInteractor: SelectionInteractor,
-            chooserRequestInteractor: ChooserRequestInteractor,
-            mimeTypeClassifier: MimeTypeClassifier,
-            // TODO: remove if possible
-            @ViewModelOwned scope: CoroutineScope,
-        ): ShareouselViewModel {
-            val keySet =
-                interactor.previews.stateIn(
-                    scope,
-                    SharingStarted.Eagerly,
-                    initialValue = null,
-                )
-            return ShareouselViewModel(
-                headline =
-                    selectionInteractor.aggregateContentType.zip(
-                        selectionInteractor.amountSelected
-                    ) { contentType, numItems ->
-                        when (contentType) {
-                            ContentType.Other -> headlineGenerator.getFilesHeadline(numItems)
-                            ContentType.Image -> headlineGenerator.getImagesHeadline(numItems)
-                            ContentType.Video -> headlineGenerator.getVideosHeadline(numItems)
+    @Provides
+    fun create(
+        interactor: SelectablePreviewsInteractor,
+        @PayloadToggle imageLoader: ImageLoader,
+        actionsInteractor: CustomActionsInteractor,
+        headlineGenerator: HeadlineGenerator,
+        selectionInteractor: SelectionInteractor,
+        chooserRequestInteractor: ChooserRequestInteractor,
+        mimeTypeClassifier: MimeTypeClassifier,
+        // TODO: remove if possible
+        @ViewModelOwned scope: CoroutineScope,
+    ): ShareouselViewModel {
+        val keySet =
+            interactor.previews.stateIn(
+                scope,
+                SharingStarted.Eagerly,
+                initialValue = null,
+            )
+        return ShareouselViewModel(
+            headline =
+                selectionInteractor.aggregateContentType.zip(selectionInteractor.amountSelected) {
+                    contentType,
+                    numItems ->
+                    when (contentType) {
+                        ContentType.Other -> headlineGenerator.getFilesHeadline(numItems)
+                        ContentType.Image -> headlineGenerator.getImagesHeadline(numItems)
+                        ContentType.Video -> headlineGenerator.getVideosHeadline(numItems)
+                    }
+                },
+            metadataText = chooserRequestInteractor.metadataText,
+            previews = keySet,
+            actions =
+                actionsInteractor.customActions.map { actions ->
+                    actions.mapIndexedNotNull { i, model ->
+                        val icon = model.icon
+                        val label = model.label
+                        if (icon == null && label.isBlank()) {
+                            null
+                        } else {
+                            ActionChipViewModel(
+                                label = label.toString(),
+                                icon = model.icon,
+                                onClicked = { model.performAction(i) },
+                            )
                         }
-                    },
-                metadataText = chooserRequestInteractor.metadataText,
-                previews = keySet,
-                actions =
-                    actionsInteractor.customActions.map { actions ->
-                        actions.mapIndexedNotNull { i, model ->
-                            val icon = model.icon
-                            val label = model.label
-                            if (icon == null && label.isBlank()) {
-                                null
-                            } else {
-                                ActionChipViewModel(
-                                    label = label.toString(),
-                                    icon = model.icon,
-                                    onClicked = { model.performAction(i) },
+                    }
+                },
+            preview = { key, previewHeight, index, previewScope ->
+                keySet.value?.maybeLoad(index)
+                val previewInteractor = interactor.preview(key)
+                val contentType =
+                    when {
+                        mimeTypeClassifier.isImageType(key.mimeType) -> ContentType.Image
+                        mimeTypeClassifier.isVideoType(key.mimeType) -> ContentType.Video
+                        else -> ContentType.Other
+                    }
+                val initialBitmapValue =
+                    key.previewUri?.let {
+                        imageLoader.getCachedBitmap(it)?.let { ValueUpdate.Value(it) }
+                    } ?: ValueUpdate.Absent
+                ShareouselPreviewViewModel(
+                    bitmapLoadState =
+                        flow {
+                                val previewWidth =
+                                    if (key.aspectRatio > 0) {
+                                            previewHeight.toFloat() / key.aspectRatio
+                                        } else {
+                                            previewHeight
+                                        }
+                                        .toInt()
+                                emit(
+                                    key.previewUri?.let {
+                                        ValueUpdate.Value(
+                                            imageLoader(it, Size(previewWidth, previewHeight))
+                                        )
+                                    } ?: ValueUpdate.Absent
                                 )
                             }
-                        }
-                    },
-                preview = { key, previewHeight, index, previewScope ->
-                    keySet.value?.maybeLoad(index)
-                    val previewInteractor = interactor.preview(key)
-                    val contentType =
-                        when {
-                            mimeTypeClassifier.isImageType(key.mimeType) -> ContentType.Image
-                            mimeTypeClassifier.isVideoType(key.mimeType) -> ContentType.Video
-                            else -> ContentType.Other
-                        }
-                    val initialBitmapValue =
-                        key.previewUri?.let {
-                            imageLoader.getCachedBitmap(it)?.let { ValueUpdate.Value(it) }
-                        } ?: ValueUpdate.Absent
-                    ShareouselPreviewViewModel(
-                        bitmapLoadState =
-                            flow {
-                                    val previewWidth =
-                                        if (key.aspectRatio > 0) {
-                                                previewHeight.toFloat() / key.aspectRatio
-                                            } else {
-                                                previewHeight
-                                            }
-                                            .toInt()
-                                    emit(
-                                        key.previewUri?.let {
-                                            ValueUpdate.Value(
-                                                imageLoader(it, Size(previewWidth, previewHeight))
-                                            )
-                                        } ?: ValueUpdate.Absent
-                                    )
-                                }
-                                .stateIn(previewScope, SharingStarted.Eagerly, initialBitmapValue),
-                        contentType = contentType,
-                        isSelected = previewInteractor.isSelected,
-                        setSelected = previewInteractor::setSelected,
-                        aspectRatio = key.aspectRatio,
-                    )
-                },
-            )
-        }
+                            .stateIn(previewScope, SharingStarted.Eagerly, initialBitmapValue),
+                    contentType = contentType,
+                    isSelected = previewInteractor.isSelected,
+                    setSelected = previewInteractor::setSelected,
+                    aspectRatio = key.aspectRatio,
+                )
+            },
+        )
     }
 }