[Storage][DataMovement] Respect transfer size options on container tr…

…ansfers and associated fixes (Azure#36487)
azure-sdk · May 25, 2023 · 8758c5e · 8758c5e
1 parent b3dc8b4
commit 8758c5e
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 24 deletions.
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/Shared/Errors.DataMovement.cs b/sdk/storage/Azure.Storage.DataMovement/src/Shared/Errors.DataMovement.cs
@@ -1,12 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+
 using System;
-using System.Collections.Generic;
 using System.IO;
-using System.Runtime.InteropServices.ComTypes;
-using System.Text;
 using Azure.Storage.DataMovement;
-using static Azure.Storage.Constants.Sas;
 
 namespace Azure.Storage
 {
@@ -92,5 +89,8 @@ public static ArgumentException MismatchResumeCreateMode(bool checkpointerValue,
             => new ArgumentException($"Mismatch Value to Resume Job: The value to overwrite / create files when they exist does not match the stored value in the transfer checkpointer. Please ensure the value passed to resume the transfer matches the value in order to prevent overwriting or failing files.\n" +
                 $"Checkpointer Value to overwrite was set to {checkpointerValue.ToString()}.\n" +
                 $"The value passed in was {passedValue.ToString()}");
+
+        public static InvalidOperationException SingleDownloadLengthMismatch(long expectedLength, long actualLength)
+            => new InvalidOperationException($"Download length {actualLength} did not match expected length {expectedLength}.");
     }
 }
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/TransferJobInternal.cs b/sdk/storage/Azure.Storage.DataMovement/src/TransferJobInternal.cs
@@ -225,6 +225,8 @@ internal TransferJobInternal(
             _sourceResourceContainer = sourceResource;
             _destinationResourceContainer = destinationResource;
             _isSingleResource = false;
+            _initialTransferSize = transferOptions?.InitialTransferSize;
+            _maximumTransferChunkSize = transferOptions?.MaximumTransferChunkSize;
         }
 
         public void Dispose()

diff --git a/sdk/storage/Azure.Storage.DataMovement/src/UriToStreamJobPart.cs b/sdk/storage/Azure.Storage.DataMovement/src/UriToStreamJobPart.cs
@@ -260,36 +260,37 @@ internal async Task LengthKnownDownloadInternal()
             {
                 await CreateZeroLengthDownload().ConfigureAwait(false);
             }
-            else if (_initialTransferSize <= totalLength)
+            // Download with a single GET
+            else if (_initialTransferSize >= totalLength)
             {
                 // To prevent requesting a range that is invalid when
                 // we already know the length we can just make one get blob request.
                 ReadStreamStorageResourceResult result = await _sourceResource.
                     ReadStreamAsync(cancellationToken: _cancellationToken)
                     .ConfigureAwait(false);
 
-                // If the initial request returned no content (i.e., a 304),
-                // we'll pass that back to the user immediately
-                long initialLength = result.Properties.ContentLength;
-                if (result == default || initialLength == 0)
+                long downloadLength = result.Properties.ContentLength;
+                // This should not occur but add a check just in case
+                if (downloadLength != totalLength)
                 {
-                    // We just need to at minimum create the file
-                    bool successfulCopy = await CopyToStreamInternal(
-                        offset: 0,
-                        sourceLength: 0,
-                        source: default,
-                        expectedLength: 0).ConfigureAwait(false);
-                    if (successfulCopy)
-                    {
-                        // Queue the work to end the download
-                        await QueueChunkToChannelAsync(
-                            async () =>
-                            await CompleteFileDownload().ConfigureAwait(false))
-                            .ConfigureAwait(false);
-                    }
-                    return;
+                    throw Errors.SingleDownloadLengthMismatch(totalLength, downloadLength);
+                }
+
+                bool successfulCopy = await CopyToStreamInternal(
+                    offset: 0,
+                    sourceLength: downloadLength,
+                    source: result.Content,
+                    expectedLength: totalLength).ConfigureAwait(false);
+                if (successfulCopy)
+                {
+                    // Queue the work to end the download
+                    await QueueChunkToChannelAsync(
+                        async () =>
+                        await CompleteFileDownload().ConfigureAwait(false))
+                        .ConfigureAwait(false);
                 }
             }
+            // Download in chunks
             else
             {
                 // Set rangeSize