Media upload cancellation (#1058)

* Initial implementation of media upload cancellation

* Add tests

* Add changelog

* Update screenshots

* Add documentation

* Fix lint issues

* Fix review comments

---------

Co-authored-by: ElementBot <benoitm+elementbot@element.io>

libraries/mediaupload/api/src/main/kotlin/io/element/android/libraries/mediaupload/api/MediaSender.kt

@@ -17,9 +17,13 @@
 package io.element.android.libraries.mediaupload.api
 
 import android.net.Uri
-import io.element.android.libraries.core.extensions.flatMap
+import io.element.android.libraries.core.extensions.flatMapCatching
 import io.element.android.libraries.matrix.api.core.ProgressCallback
+import io.element.android.libraries.matrix.api.media.MediaUploadHandler
 import io.element.android.libraries.matrix.api.room.MatrixRoom
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.Job
+import java.util.concurrent.ConcurrentHashMap
 import javax.inject.Inject
 
 class MediaSender @Inject constructor(
@@ -27,6 +31,9 @@ class MediaSender @Inject constructor(
     private val room: MatrixRoom,
 ) {
 
+    private val ongoingUploadJobs = ConcurrentHashMap<Job.Key, MediaUploadHandler>()
+    val hasOngoingMediaUploads get() = ongoingUploadJobs.isNotEmpty()
+
     suspend fun sendMedia(
         uri: Uri,
         mimeType: String,
@@ -40,16 +47,25 @@ class MediaSender @Inject constructor(
                 deleteOriginal = true,
                 compressIfPossible = compressIfPossible
             )
-            .flatMap { info ->
+            .flatMapCatching { info ->
                 room.sendMedia(info, progressCallback)
             }
+            .onFailure { error ->
+                val job = ongoingUploadJobs.remove(Job)
+                if (error !is CancellationException) {
+                    job?.cancel()
+                }
+            }
+            .onSuccess {
+                ongoingUploadJobs.remove(Job)
+            }
     }
 
     private suspend fun MatrixRoom.sendMedia(
         uploadInfo: MediaUploadInfo,
-        progressCallback: ProgressCallback?
+        progressCallback: ProgressCallback?,
     ): Result<Unit> {
-        return when (uploadInfo) {
+        val handler = when (uploadInfo) {
             is MediaUploadInfo.Image -> {
                 sendImage(
                     file = uploadInfo.file,
@@ -83,5 +99,11 @@ class MediaSender @Inject constructor(
                 )
             }
         }
+
+        return handler
+            .flatMapCatching { uploadHandler ->
+                ongoingUploadJobs[Job] = uploadHandler
+                uploadHandler.await()
+            }
     }
 }

libraries/mediaupload/api/src/test/kotlin/io/element/android/libraries/mediaupload/api/MediaSenderTests.kt

@@ -0,0 +1,116 @@
+/*
+ * Copyright (c) 2023 New Vector Ltd
+ *
+ * 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 io.element.android.libraries.mediaupload.api
+
+import android.net.Uri
+import com.google.common.truth.Truth.assertThat
+import io.element.android.libraries.matrix.api.room.MatrixRoom
+import io.element.android.libraries.matrix.test.room.FakeMatrixRoom
+import io.element.android.libraries.mediaupload.test.FakeMediaPreProcessor
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.advanceTimeBy
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@RunWith(RobolectricTestRunner::class)
+class MediaSenderTests {
+
+    @Test
+    fun `given an attachment when sending it the preprocessor always runs`() = runTest {
+        val preProcessor = FakeMediaPreProcessor()
+        val sender = aMediaSender(preProcessor)
+
+        val uri = Uri.parse("content://image.jpg")
+        sender.sendMedia(uri = uri, mimeType = "image/jpeg", compressIfPossible = true)
+
+        assertThat(preProcessor.processCallCount).isEqualTo(1)
+    }
+
+    @Test
+    fun `given an attachment when sending it the MatrixRoom will call sendMedia`() = runTest {
+        val room = FakeMatrixRoom()
+        val sender = aMediaSender(room = room)
+
+        val uri = Uri.parse("content://image.jpg")
+        sender.sendMedia(uri = uri, mimeType = "image/jpeg", compressIfPossible = true)
+
+        assertThat(room.sendMediaCount).isEqualTo(1)
+    }
+
+    @Test
+    fun `given a failure in the preprocessor when sending the whole process fails`() = runTest {
+        val preProcessor = FakeMediaPreProcessor().apply {
+            givenResult(Result.failure(Exception()))
+        }
+        val sender = aMediaSender(preProcessor)
+
+        val uri = Uri.parse("content://image.jpg")
+        val result = sender.sendMedia(uri = uri, mimeType = "image/jpeg", compressIfPossible = true)
+
+        assertThat(result.exceptionOrNull()).isNotNull()
+    }
+
+    @Test
+    fun `given a failure in the media upload when sending the whole process fails`() = runTest {
+        val room = FakeMatrixRoom().apply {
+            givenSendMediaResult(Result.failure(Exception()))
+        }
+        val sender = aMediaSender(room = room)
+
+        val uri = Uri.parse("content://image.jpg")
+        val result = sender.sendMedia(uri = uri, mimeType = "image/jpeg", compressIfPossible = true)
+
+        assertThat(result.exceptionOrNull()).isNotNull()
+    }
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    @Test
+    fun `given a cancellation in the media upload when sending the job is cancelled`() = runTest(StandardTestDispatcher()) {
+        val room = FakeMatrixRoom()
+        val sender = aMediaSender(room = room)
+        val sendJob = launch {
+            val uri = Uri.parse("content://image.jpg")
+            sender.sendMedia(uri = uri, mimeType = "image/jpeg", compressIfPossible = true)
+        }
+        // Wait until several internal tasks run and the file is being uploaded
+        advanceTimeBy(3L)
+
+        // Assert the file is being uploaded
+        assertThat(sender.hasOngoingMediaUploads).isTrue()
+
+        // Cancel the coroutine
+        sendJob.cancel()
+
+        // Wait for the coroutine cleanup to happen
+        advanceTimeBy(1L)
+
+        // Assert the file is not being uploaded anymore
+        assertThat(sender.hasOngoingMediaUploads).isFalse()
+    }
+
+    private fun aMediaSender(
+        preProcessor: MediaPreProcessor = FakeMediaPreProcessor(),
+        room: MatrixRoom = FakeMatrixRoom(),
+    ) = MediaSender(
+        preProcessor,
+        room,
+    )
+}