From bbb929ff3ed7e548bcc7ee38e80fa2f1b06dacc1 Mon Sep 17 00:00:00 2001
From: Tyler-Larkin <tyllark@amazon.com>
Date: Thu, 21 Nov 2024 10:19:49 -0800
Subject: [PATCH 1/2] fix(storage): Only allow 1 batch to run at a time

---
 .../src/storage_s3_service/service/task/s3_upload_task.dart   | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/packages/storage/amplify_storage_s3_dart/lib/src/storage_s3_service/service/task/s3_upload_task.dart b/packages/storage/amplify_storage_s3_dart/lib/src/storage_s3_service/service/task/s3_upload_task.dart
index 55f5091165..dbecca6d44 100644
--- a/packages/storage/amplify_storage_s3_dart/lib/src/storage_s3_service/service/task/s3_upload_task.dart
+++ b/packages/storage/amplify_storage_s3_dart/lib/src/storage_s3_service/service/task/s3_upload_task.dart
@@ -547,11 +547,15 @@ class S3UploadTask {
     }
   }
 
+  bool _isNextBatchWaiting = false;
   Future<void> _startNextUploadPartsBatch({
     bool resumingFromPause = false,
   }) async {
     // await for previous batching to complete (if any)
+    if (_isNextBatchWaiting) return;
+    _isNextBatchWaiting = true;
     await _uploadPartBatchingCompleted;
+    _isNextBatchWaiting = false;
 
     if (_state != StorageTransferState.inProgress) {
       return;

From dcf4ab491d1f944fc880e0a0835d7eadbf6da828 Mon Sep 17 00:00:00 2001
From: Tyler-Larkin <tyllark@amazon.com>
Date: Fri, 22 Nov 2024 15:34:24 -0800
Subject: [PATCH 2/2] chore(storage): Added async multipart read test

---
 .../task/s3_upload_task_test.dart             | 103 ++++++++++++++++++
 1 file changed, 103 insertions(+)

diff --git a/packages/storage/amplify_storage_s3_dart/test/storage_s3_service/task/s3_upload_task_test.dart b/packages/storage/amplify_storage_s3_dart/test/storage_s3_service/task/s3_upload_task_test.dart
index 094aef27c4..51d373d0c5 100644
--- a/packages/storage/amplify_storage_s3_dart/test/storage_s3_service/task/s3_upload_task_test.dart
+++ b/packages/storage/amplify_storage_s3_dart/test/storage_s3_service/task/s3_upload_task_test.dart
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import 'dart:async';
+import 'dart:math';
 import 'dart:typed_data';
 
 import 'package:amplify_core/amplify_core.dart';
@@ -1238,6 +1239,108 @@ void main() {
         expect(finalState, StorageTransferState.failure);
       });
 
+      test('should handle async gaps when reading from Multipart file',
+          () async {
+        late StorageTransferState finalState;
+
+        //completeMultipartUploadSmithyOperation
+        final testCompleteMultipartUploadOutput =
+            s3.CompleteMultipartUploadOutput();
+        final completeMultipartUploadSmithyOperation =
+            MockSmithyOperation<s3.CompleteMultipartUploadOutput>();
+        when(
+          () => completeMultipartUploadSmithyOperation.result,
+        ).thenAnswer(
+          (_) async => testCompleteMultipartUploadOutput,
+        );
+
+        //uploadPartSmithyOperation
+        final testUploadPartOutput = s3.UploadPartOutput(eTag: 'eTag-part-1');
+        final uploadPartSmithyOperation =
+            MockSmithyOperation<s3.UploadPartOutput>();
+        when(
+          () => uploadPartSmithyOperation.result,
+        ).thenAnswer(
+          (_) async => testUploadPartOutput,
+        );
+
+        //createMultipartUploadSmithyOperation
+        final testCreateMultipartUploadOutput = s3.CreateMultipartUploadOutput(
+          uploadId: 'uploadId', // response should always contain valid uploadId
+        );
+        final createMultipartUploadSmithyOperation =
+            MockSmithyOperation<s3.CreateMultipartUploadOutput>();
+        when(
+          () => createMultipartUploadSmithyOperation.result,
+        ).thenAnswer(
+          (_) async => testCreateMultipartUploadOutput,
+        );
+
+        //s3Client
+        when(
+          () => s3Client.completeMultipartUpload(any()),
+        ).thenAnswer((_) => completeMultipartUploadSmithyOperation);
+        when(
+          () => s3Client.uploadPart(
+            any(),
+            s3ClientConfig: any(named: 's3ClientConfig'),
+          ),
+        ).thenAnswer(
+          (_) => uploadPartSmithyOperation,
+        );
+        when(
+          () => s3Client.createMultipartUpload(any()),
+        ).thenAnswer(
+          (_) => createMultipartUploadSmithyOperation,
+        );
+
+        //transferDatabase
+        when(
+          () => transferDatabase.insertTransferRecord(any<TransferRecord>()),
+        ).thenAnswer(
+          (_) async => '1',
+        );
+        when(
+          () => transferDatabase.deleteTransferRecords(any()),
+        ).thenAnswer(
+          (_) async => 1,
+        );
+
+        final bytes = List<int>.filled(
+          (32 * pow(2, 20)).toInt(),
+          0,
+        );
+        final mockFile = AWSFile.fromStream(
+          Stream.value(bytes),
+          size: bytes.length,
+          contentType: 'image/jpeg',
+        );
+
+        final uploadTask = S3UploadTask.fromAWSFile(
+          mockFile,
+          s3Client: s3Client,
+          defaultS3ClientConfig: defaultS3ClientConfig,
+          pathResolver: pathResolver,
+          bucket: testBucket,
+          path: const StoragePath.fromString(testKey),
+          options: testUploadDataOptions,
+          logger: logger,
+          transferDatabase: transferDatabase,
+          onProgress: (progress) {
+            finalState = progress.state;
+          },
+        );
+
+        unawaited(uploadTask.start());
+
+        await uploadTask.result;
+
+        expect(
+          finalState,
+          StorageTransferState.success,
+        );
+      });
+
       test(
           'should complete with StorageAccessDeniedException when CreateMultipartUploadRequest'
           ' returned UnknownSmithyHttpException with status code 403',