fix(storage): Download Accuracy Fix (#3156)

Co-authored-by: Matt Creaser <[email protected]>

Author: Tyler Roach

aws-storage-s3/src/main/java/com/amplifyframework/storage/s3/TransferOperations.kt

@@ -38,6 +38,7 @@ import com.amplifyframework.storage.s3.transfer.worker.InitiateMultiPartUploadTr
 import com.amplifyframework.storage.s3.transfer.worker.PartUploadTransferWorker
 import com.amplifyframework.storage.s3.transfer.worker.RouterWorker
 import com.amplifyframework.storage.s3.transfer.worker.SinglePartUploadWorker
+import java.util.concurrent.TimeUnit
 
 internal object TransferOperations {
 
@@ -84,7 +85,13 @@ internal object TransferOperations {
     ): Boolean {
         if (TransferState.isStarted(transferRecord.state) && !TransferState.isInTerminalState(transferRecord.state)) {
             transferStatusUpdater.updateTransferState(transferRecord.id, TransferState.PENDING_PAUSE)
-            workManager.cancelUniqueWork(transferRecord.id.toString())
+            try {
+                workManager.cancelUniqueWork(
+                    transferRecord.id.toString()
+                ).result.get(1, TimeUnit.SECONDS)
+            } catch (_: Exception) {
+                // do nothing
+            }
             return true
         }
         return false

aws-storage-s3/src/main/java/com/amplifyframework/storage/s3/transfer/ClearableBufferedOutputStream.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License").
+ *  You may not use this file except in compliance with the License.
+ *  A copy of the License is located at
+ *
+ *   http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file 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.amplifyframework.storage.s3.transfer
+
+import java.io.BufferedOutputStream
+import java.io.OutputStream
+
+/**
+ * BufferedOutputStream that allows clearing the buffer without flushing. This may especially be important when
+ * attempting to close the BufferedOutputStream without flushing the buffer.
+ */
+internal class ClearableBufferedOutputStream(
+    out: OutputStream,
+    size: Int = 8192
+) : BufferedOutputStream(out, size) {
+
+    fun clear() {
+        count = 0
+        buf.fill(0)
+    }
+}

aws-storage-s3/src/main/java/com/amplifyframework/storage/s3/transfer/worker/DownloadWorker.kt

@@ -24,16 +24,15 @@ import aws.smithy.kotlin.runtime.content.ByteStream
 import aws.smithy.kotlin.runtime.content.writeToFile
 import aws.smithy.kotlin.runtime.io.SdkSource
 import aws.smithy.kotlin.runtime.io.buffer
+import com.amplifyframework.storage.s3.transfer.ClearableBufferedOutputStream
 import com.amplifyframework.storage.s3.transfer.DownloadProgressListener
 import com.amplifyframework.storage.s3.transfer.DownloadProgressListenerInterceptor
 import com.amplifyframework.storage.s3.transfer.StorageTransferClientProvider
 import com.amplifyframework.storage.s3.transfer.TransferDB
 import com.amplifyframework.storage.s3.transfer.TransferStatusUpdater
-import java.io.BufferedOutputStream
 import java.io.File
 import java.io.FileOutputStream
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.withContext
 
@@ -101,9 +100,11 @@ internal class DownloadWorker(
                     val append = file.length() > 0
                     val fileOutputStream = FileOutputStream(file, append)
                     var totalRead = 0L
-                    BufferedOutputStream(fileOutputStream).use { fileOutput ->
+                    // use ensures the underlying source is closed. In this case, a BufferedOutputStream. By default,
+                    // a bos flushes on close. We may not want this behavior, so we use ClearableBufferedOutputStream.
+                    ClearableBufferedOutputStream(fileOutputStream).use { fileOutput ->
                         val copied = 0L
-                        while (currentCoroutineContext().isActive) {
+                        while (isActive) {
                             val remaining = limit - copied
                             if (remaining == 0L) break
                             val readBytes =
@@ -112,10 +113,24 @@ internal class DownloadWorker(
                             if (readBytes > 0) {
                                 totalRead += readBytes
                             }
-                            fileOutput.write(buffer, 0, readBytes)
+                            if (isActive) {
+                                // Double check to make sure that we are still active before writing to buffer
+                                fileOutput.write(buffer, 0, readBytes)
+                            } else {
+                                // If we are no longer active, clear the buffer so that no more data is written to the
+                                // file. A resume operation may have already started, and it resumes based on the
+                                // file size at its start. A flush here could result in duplicating file data
+                                fileOutput.clear()
+                            }
                         }
-                        if (sourceStream.buffer().exhausted()) {
+                        if (sourceStream.buffer().exhausted() && isActive) {
+                            // Double check to make sure that we are still active before flushing to buffer
                             fileOutput.flush()
+                        } else {
+                            // If we are no longer active, clear the buffer so that no more data is written to the
+                            // file. A resume operation may have already started, and it resumes based on the
+                            // file size at its start. A flush here could result in duplicating file data
+                            fileOutput.clear()
                         }
                     }
                 }