[Storage][DataMovement] Refactor some internal transfer logic (Azure#36367)

jalauzon-msft · web-flow · commit bb676d880399 · 2023-05-22T14:38:03.000-07:00
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/JobPartInternal.cs b/sdk/storage/Azure.Storage.DataMovement/src/JobPartInternal.cs
@@ -254,16 +254,6 @@ await QueueChunk(
         /// <returns>The task to wait until the cancellation has been triggered.</returns>
         internal async Task TriggerCancellationAsync()
         {
-            // If stop on failure specified, cancel entire job.
-            if (_errorHandling == ErrorHandlingOptions.StopOnAllFailures)
-            {
-                if (!_cancellationToken.IsCancellationRequested)
-                {
-                    _dataTransfer._state.TriggerCancellation();
-                }
-                _dataTransfer._state.ResetTransferredBytes();
-            }
-
             // Set the status to Pause/CancellationInProgress
             if (StorageTransferStatus.PauseInProgress == _dataTransfer.TransferStatus)
             {
@@ -276,6 +266,7 @@ internal async Task TriggerCancellationAsync()
                 // It's a cancellation if a pause wasn't called.
                 await OnTransferStatusChanged(StorageTransferStatus.CancellationInProgress).ConfigureAwait(false);
             }
+            await CleanupAbortedJobPartAsync().ConfigureAwait(false);
         }
 
         /// <summary>
@@ -300,11 +291,7 @@ internal async Task OnTransferStatusChanged(StorageTransferStatus transferStatus
                 {
                     await InvokeSingleCompletedArg().ConfigureAwait(false);
                 }
-                if (JobPartStatus == StorageTransferStatus.Paused ||
-                    JobPartStatus == StorageTransferStatus.CompletedWithFailedTransfers)
-                {
-                    await CleanupAbortedJobPartAsync().ConfigureAwait(false);
-                }
+
                 // Set the status in the checkpointer
                 await SetCheckpointerStatus(transferStatus).ConfigureAwait(false);
 
@@ -510,7 +497,7 @@ internal static long ParseRangeTotalLength(string range)
             long absolutePosition = blockSize;
             long blockLength = acceptableBlockSize;
 
-            // TODO: divide up paritions based on how much array pool is left
+            // TODO: divide up partitions based on how much array pool is left
             while (absolutePosition < streamLength)
             {
                 // Return based on the size of the stream divided up by the acceptable blocksize.
@@ -526,11 +513,11 @@ internal async Task CheckAndUpdateCancellationStatusAsync()
         {
             if (_chunkTasks.All((Task task) => (task.IsCompleted)))
             {
-                if (_dataTransfer.TransferStatus == StorageTransferStatus.PauseInProgress)
+                if (JobPartStatus == StorageTransferStatus.PauseInProgress)
                 {
                     await OnTransferStatusChanged(StorageTransferStatus.Paused).ConfigureAwait(false);
                 }
-                else if (_dataTransfer.TransferStatus == StorageTransferStatus.CancellationInProgress)
+                else if (JobPartStatus == StorageTransferStatus.CancellationInProgress)
                 {
                     await OnTransferStatusChanged(StorageTransferStatus.CompletedWithFailedTransfers).ConfigureAwait(false);
                 }
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/TransferJobInternal.cs b/sdk/storage/Azure.Storage.DataMovement/src/TransferJobInternal.cs
@@ -260,21 +260,14 @@ public void DisposeHandlers()
         /// <returns>The task to wait until the cancellation has been triggered.</returns>
         public async Task TriggerJobCancellationAsync()
         {
-            DisposeHandlers();
             if (!_dataTransfer._state.CancellationTokenSource.IsCancellationRequested)
             {
                 _dataTransfer._state.TriggerCancellation();
             }
-            // Set the status to Pause/CancellationInProgress
-            if (StorageTransferStatus.PauseInProgress == _dataTransfer.TransferStatus)
-            {
-                // It's possible that the status hasn't propagated down to the job part
-                // status yet here since we pause from the data transfer object.
-                await OnJobStatusChangedAsync(StorageTransferStatus.PauseInProgress).ConfigureAwait(false);
-            }
-            else
+
+            // Set CancellationInProgress if we are not already pausing
+            if (StorageTransferStatus.PauseInProgress != _dataTransfer.TransferStatus)
             {
-                // It's a cancellation if a pause wasn't called.
                 await OnJobStatusChangedAsync(StorageTransferStatus.CancellationInProgress).ConfigureAwait(false);
             }
         }
@@ -315,30 +308,33 @@ await TransferFailedEventHandler.Invoke(new TransferFailedEventArgs(
         /// </summary>
         public async Task JobPartEvent(TransferStatusEventArgs args)
         {
-            StorageTransferStatus status = _dataTransfer._state.GetTransferStatus();
-            if ((args.StorageTransferStatus == StorageTransferStatus.Paused ||
-                 args.StorageTransferStatus == StorageTransferStatus.Completed ||
-                 args.StorageTransferStatus == StorageTransferStatus.CompletedWithSkippedTransfers ||
-                 args.StorageTransferStatus == StorageTransferStatus.CompletedWithFailedTransfers)
-                && (status == StorageTransferStatus.Queued ||
-                    status == StorageTransferStatus.InProgress ||
-                    status == StorageTransferStatus.PauseInProgress ||
-                    status == StorageTransferStatus.CancellationInProgress))
+            StorageTransferStatus jobPartStatus = args.StorageTransferStatus;
+            StorageTransferStatus jobStatus = _dataTransfer._state.GetTransferStatus();
+
+            // Cancel the entire job if one job part fails and StopOnFailure is set
+            if (_errorHandling == ErrorHandlingOptions.StopOnAllFailures &&
+                jobPartStatus == StorageTransferStatus.CompletedWithFailedTransfers &&
+                jobStatus != StorageTransferStatus.CompletedWithFailedTransfers &&
+                jobStatus != StorageTransferStatus.CompletedWithSkippedTransfers &&
+                jobStatus != StorageTransferStatus.Completed)
+            {
+                await TriggerJobCancellationAsync().ConfigureAwait(false);
+            }
+
+            if ((jobPartStatus == StorageTransferStatus.Paused ||
+                 jobPartStatus == StorageTransferStatus.Completed ||
+                 jobPartStatus == StorageTransferStatus.CompletedWithSkippedTransfers ||
+                 jobPartStatus == StorageTransferStatus.CompletedWithFailedTransfers)
+                && (jobStatus == StorageTransferStatus.Queued ||
+                    jobStatus == StorageTransferStatus.InProgress ||
+                    jobStatus == StorageTransferStatus.PauseInProgress ||
+                    jobStatus == StorageTransferStatus.CancellationInProgress))
             {
                 if (_enumerationComplete)
                 {
                     await CheckAndUpdateStatusAsync().ConfigureAwait(false);
                 }
             }
-            else if (args.StorageTransferStatus == StorageTransferStatus.Paused &&
-                    status == StorageTransferStatus.Paused)
-            {
-                await OnJobStatusChangedAsync(StorageTransferStatus.Paused).ConfigureAwait(false);
-            }
-            else if (args.StorageTransferStatus > status)
-            {
-                await OnJobStatusChangedAsync(args.StorageTransferStatus).ConfigureAwait(false);
-            }
         }
 
         public async Task OnJobStatusChangedAsync(StorageTransferStatus status)
@@ -350,6 +346,15 @@ public async Task OnJobStatusChangedAsync(StorageTransferStatus status)
             }
             if (statusChanged)
             {
+                // If we are in a final state, dispose the JobPartEvent handlers
+                if (status == StorageTransferStatus.Completed ||
+                    status == StorageTransferStatus.CompletedWithSkippedTransfers ||
+                    status == StorageTransferStatus.CompletedWithFailedTransfers ||
+                    status == StorageTransferStatus.Paused)
+                {
+                    DisposeHandlers();
+                }
+
                 if (TransferStatusEventHandler != null)
                 {
                     await TransferStatusEventHandler.Invoke(
diff --git a/sdk/storage/Azure.Storage.DataMovement/tests/PauseResumeTransferTests.cs b/sdk/storage/Azure.Storage.DataMovement/tests/PauseResumeTransferTests.cs
@@ -447,7 +447,7 @@ public async Task PauseThenResumeTransferAsync(TransferType transferType)
             {
                 ResumeFromCheckpointId = transfer.Id
             };
-            TestEventsRaised resumeFailureHolder = new TestEventsRaised(resumeOptions);
+            TestEventsRaised testEventRaised2 = new TestEventsRaised(resumeOptions);
             DataTransfer resumeTransfer = await CreateSingleLongTransferAsync(
                 manager: transferManager,
                 sourceResource: sourceResource,
@@ -458,7 +458,7 @@ public async Task PauseThenResumeTransferAsync(TransferType transferType)
             await resumeTransfer.AwaitCompletion(waitTransferCompletion.Token);
 
             // Assert
-            resumeFailureHolder.AssertPausedCheck();
+            testEventRaised2.AssertContainerCompletedCheck(1);
             Assert.AreEqual(StorageTransferStatus.Completed, resumeTransfer.TransferStatus);
             Assert.IsTrue(resumeTransfer.HasCompleted);
 
@@ -809,18 +809,19 @@ public async Task PauseThenResumeTransferAsync_Directory(TransferType transferTy
             {
                 ResumeFromCheckpointId = transfer.Id
             };
-            TestEventsRaised resumeFailureHolder = new TestEventsRaised(resumeOptions);
+            TestEventsRaised testEventsRaised2 = new TestEventsRaised(resumeOptions);
             DataTransfer resumeTransfer = await CreateDirectoryLongTransferAsync(
                 manager: transferManager,
                 sourceResource: sourceResource,
                 destinationResource: destinationResource,
-                transferOptions: resumeOptions);
+                transferOptions: resumeOptions,
+                transferCount: 100);
 
             CancellationTokenSource waitTransferCompletion = new CancellationTokenSource(TimeSpan.FromSeconds(600));
             await resumeTransfer.AwaitCompletion(waitTransferCompletion.Token);
 
             // Assert
-            resumeFailureHolder.AssertPausedCheck();
+            testEventsRaised2.AssertContainerCompletedCheck(100);
             Assert.AreEqual(StorageTransferStatus.Completed, resumeTransfer.TransferStatus);
             Assert.IsTrue(resumeTransfer.HasCompleted);
 
diff --git a/sdk/storage/Azure.Storage.DataMovement/tests/StartTransferDownloadDirectoryTests.cs b/sdk/storage/Azure.Storage.DataMovement/tests/StartTransferDownloadDirectoryTests.cs
@@ -329,6 +329,79 @@ await DownloadBlobDirectoryAndVerify(
                 destinationFolder).ConfigureAwait(false);
         }
 
+        [Test]
+        [LiveOnly] // https://github.com/Azure/azure-sdk-for-net/issues/33082
+        public async Task DownloadDirectoryAsync_SmallChunks_ManyFiles()
+        {
+            // Arrange
+            int blobSize = 2 * Constants.KB;
+            await using DisposingBlobContainer test = await GetTestContainerAsync();
+
+            using DisposingLocalDirectory testDirectory = DisposingLocalDirectory.GetTestDirectory();
+            string tempFolder = CreateRandomDirectory(testDirectory.DirectoryPath);
+            string blobDirectoryName = "foo";
+            string fullSourceFolderPath = CreateRandomDirectory(tempFolder, blobDirectoryName);
+
+            List<string> blobNames = new List<string>();
+            string blobName = Path.Combine(blobDirectoryName, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+            blobName = Path.Combine(blobDirectoryName, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+            blobName = Path.Combine(blobDirectoryName, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+            blobName = Path.Combine(blobDirectoryName, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+            blobName = Path.Combine(blobDirectoryName, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+
+            string subDir1 = CreateRandomDirectory(fullSourceFolderPath, "bar").Substring(fullSourceFolderPath.Length + 1);
+            blobName = Path.Combine(blobDirectoryName, subDir1, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+            blobName = Path.Combine(blobDirectoryName, subDir1, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+            blobName = Path.Combine(blobDirectoryName, subDir1, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+            string subDir2 = CreateRandomDirectory(fullSourceFolderPath, "rul").Substring(fullSourceFolderPath.Length + 1);
+            blobName = Path.Combine(blobDirectoryName, subDir2, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+            blobName = Path.Combine(blobDirectoryName, subDir2, GetNewBlobName());
+            await CreateBlockBlobAndSourceFile(test.Container, tempFolder, blobName, blobSize);
+            blobNames.Add(blobName);
+
+            using DisposingLocalDirectory destinationFolder = DisposingLocalDirectory.GetTestDirectory();
+            string sourceBlobPrefix = fullSourceFolderPath.Substring(tempFolder.Length + 1);
+
+            TransferManagerOptions transferManagerOptions = new TransferManagerOptions()
+            {
+                ErrorHandling = ErrorHandlingOptions.StopOnAllFailures,
+                MaximumConcurrency = 3
+            };
+            TransferOptions options = new TransferOptions()
+            {
+                InitialTransferSize = 512,
+                MaximumTransferChunkSize = 512
+            };
+
+            // Act / Assert
+            await DownloadBlobDirectoryAndVerify(
+                sourceContainer: test.Container,
+                sourceBlobPrefix: sourceBlobPrefix,
+                sourceFilePrefix: fullSourceFolderPath,
+                blobNames,
+                destinationFolder.DirectoryPath,
+                transferManagerOptions: transferManagerOptions,
+                options: options).ConfigureAwait(false);
+        }
+
         [Test]
         [LiveOnly] // https://github.com/Azure/azure-sdk-for-net/issues/33082
         public async Task DownloadDirectoryAsync_Root()
@@ -436,7 +509,7 @@ public async Task StartTransfer_AwaitCompletion()
                 options: options);
 
             // Act
-            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10));
             await transfer.AwaitCompletion(cancellationTokenSource.Token).ConfigureAwait(false);
 
             // Assert
@@ -446,7 +519,6 @@ public async Task StartTransfer_AwaitCompletion()
             testEventsRaised.AssertContainerCompletedCheck(4);
         }
 
-        //[Ignore("https://github.com/Azure/azure-sdk-for-net/issues/35209")]
         [Test]
         [LiveOnly] // https://github.com/Azure/azure-sdk-for-net/issues/33082
         public async Task StartTransfer_AwaitCompletion_Failed()
@@ -473,7 +545,7 @@ public async Task StartTransfer_AwaitCompletion_Failed()
                 options: options);
 
             // Act
-            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10));
             await transfer.AwaitCompletion(cancellationTokenSource.Token).ConfigureAwait(false);
 
             // Assert
@@ -484,7 +556,6 @@ public async Task StartTransfer_AwaitCompletion_Failed()
             testEventsRaised.AssertContainerCompletedWithFailedCheck(1);
         }
 
-        //[Ignore("https://github.com/Azure/azure-sdk-for-net/issues/35209")]
         [Test]
         [LiveOnly] // https://github.com/Azure/azure-sdk-for-net/issues/33082
         public async Task StartTransfer_AwaitCompletion_Skipped()
@@ -512,7 +583,7 @@ public async Task StartTransfer_AwaitCompletion_Skipped()
                 options: options);
 
             // Act
-            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10));
             await transfer.AwaitCompletion(cancellationTokenSource.Token).ConfigureAwait(false);
 
             // Assert
@@ -542,7 +613,7 @@ public async Task StartTransfer_EnsureCompleted()
                     options: options);
 
             // Act
-            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10));
             transfer.EnsureCompleted(cancellationTokenSource.Token);
 
             // Assert
@@ -552,7 +623,6 @@ public async Task StartTransfer_EnsureCompleted()
             testEventsRaised.AssertContainerCompletedCheck(4);
         }
 
-        //[Ignore("https://github.com/Azure/azure-sdk-for-net/issues/35209")]
         [Test]
         [LiveOnly] // https://github.com/Azure/azure-sdk-for-net/issues/33082
         public async Task StartTransfer_EnsureCompleted_Failed()
@@ -579,7 +649,7 @@ public async Task StartTransfer_EnsureCompleted_Failed()
                 options: options);
 
             // Act
-            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10));
             transfer.EnsureCompleted(cancellationTokenSource.Token);
 
             // Assert
@@ -590,7 +660,6 @@ public async Task StartTransfer_EnsureCompleted_Failed()
             testEventsRaised.AssertContainerCompletedWithFailedCheck(1);
         }
 
-        //[Ignore("https://github.com/Azure/azure-sdk-for-net/issues/35209")]
         [Test]
         [LiveOnly] // https://github.com/Azure/azure-sdk-for-net/issues/33082
         public async Task StartTransfer_EnsureCompleted_Skipped()
@@ -618,7 +687,7 @@ public async Task StartTransfer_EnsureCompleted_Skipped()
                 options: options);
 
             // Act
-            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10));
             transfer.EnsureCompleted(cancellationTokenSource.Token);
 
             // Assert
diff --git a/sdk/storage/Azure.Storage.DataMovement/tests/StartTransferUploadDirectoryTests.cs b/sdk/storage/Azure.Storage.DataMovement/tests/StartTransferUploadDirectoryTests.cs
diff --git a/sdk/storage/Azure.Storage.DataMovement/tests/TestEventsRaised.cs b/sdk/storage/Azure.Storage.DataMovement/tests/TestEventsRaised.cs

Original file line number	Diff line number	Diff line change
`@@ -254,16 +254,6 @@ await QueueChunk(`
`254`	`254`	`/// <returns>The task to wait until the cancellation has been triggered.</returns>`
`255`	`255`	`internal async Task TriggerCancellationAsync()`
`256`	`256`	`{`
`257`		`- // If stop on failure specified, cancel entire job.`
`258`		`- if (_errorHandling == ErrorHandlingOptions.StopOnAllFailures)`
`259`		`- {`
`260`		`- if (!_cancellationToken.IsCancellationRequested)`
`261`		`- {`
`262`		`- _dataTransfer._state.TriggerCancellation();`
`263`		`- }`
`264`		`- _dataTransfer._state.ResetTransferredBytes();`
`265`		`- }`
`266`		`-`
`267`	`257`	`// Set the status to Pause/CancellationInProgress`
`268`	`258`	`if (StorageTransferStatus.PauseInProgress == _dataTransfer.TransferStatus)`
`269`	`259`	`{`
`@@ -276,6 +266,7 @@ internal async Task TriggerCancellationAsync()`
`276`	`266`	`// It's a cancellation if a pause wasn't called.`
`277`	`267`	`await OnTransferStatusChanged(StorageTransferStatus.CancellationInProgress).ConfigureAwait(false);`
`278`	`268`	`}`
	`269`	`+ await CleanupAbortedJobPartAsync().ConfigureAwait(false);`
`279`	`270`	`}`
`280`	`271`
`281`	`272`	`/// <summary>`
`@@ -300,11 +291,7 @@ internal async Task OnTransferStatusChanged(StorageTransferStatus transferStatus`
`300`	`291`	`{`
`301`	`292`	`await InvokeSingleCompletedArg().ConfigureAwait(false);`
`302`	`293`	`}`
`303`		`- if (JobPartStatus == StorageTransferStatus.Paused \|\|`
`304`		`- JobPartStatus == StorageTransferStatus.CompletedWithFailedTransfers)`
`305`		`- {`
`306`		`- await CleanupAbortedJobPartAsync().ConfigureAwait(false);`
`307`		`- }`
	`294`	`+`
`308`	`295`	`// Set the status in the checkpointer`
`309`	`296`	`await SetCheckpointerStatus(transferStatus).ConfigureAwait(false);`
`310`	`297`
`@@ -510,7 +497,7 @@ internal static long ParseRangeTotalLength(string range)`
`510`	`497`	`long absolutePosition = blockSize;`
`511`	`498`	`long blockLength = acceptableBlockSize;`
`512`	`499`
`513`		`- // TODO: divide up paritions based on how much array pool is left`
	`500`	`+ // TODO: divide up partitions based on how much array pool is left`
`514`	`501`	`while (absolutePosition < streamLength)`
`515`	`502`	`{`
`516`	`503`	`// Return based on the size of the stream divided up by the acceptable blocksize.`
`@@ -526,11 +513,11 @@ internal async Task CheckAndUpdateCancellationStatusAsync()`
`526`	`513`	`{`
`527`	`514`	`if (_chunkTasks.All((Task task) => (task.IsCompleted)))`
`528`	`515`	`{`
`529`		`- if (_dataTransfer.TransferStatus == StorageTransferStatus.PauseInProgress)`
	`516`	`+ if (JobPartStatus == StorageTransferStatus.PauseInProgress)`
`530`	`517`	`{`
`531`	`518`	`await OnTransferStatusChanged(StorageTransferStatus.Paused).ConfigureAwait(false);`
`532`	`519`	`}`
`533`		`- else if (_dataTransfer.TransferStatus == StorageTransferStatus.CancellationInProgress)`
	`520`	`+ else if (JobPartStatus == StorageTransferStatus.CancellationInProgress)`
`534`	`521`	`{`
`535`	`522`	`await OnTransferStatusChanged(StorageTransferStatus.CompletedWithFailedTransfers).ConfigureAwait(false);`
`536`	`523`	`}`