Dispose fixes on download (#46592)

* dispose fixes * changelog * grammar * feedback
Azure · Oct 28, 2024 · 597a9a7 · 597a9a7
1 parent 4160653
commit 597a9a7
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/sdk/storage/Azure.Storage.Blobs/CHANGELOG.md b/sdk/storage/Azure.Storage.Blobs/CHANGELOG.md
@@ -7,6 +7,8 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+- Fixed bug where network download streams were not properly disposed.
+- Fixed bug where DownloadToAsync() did not dispose all its network streams on error in some cases.
 
 ### Other Changes
 

diff --git a/sdk/storage/Azure.Storage.Blobs/src/BlobBaseClient.cs b/sdk/storage/Azure.Storage.Blobs/src/BlobBaseClient.cs
@@ -1572,7 +1572,7 @@ internal virtual async ValueTask<Response<BlobDownloadStreamingResult>> Download
                         ClientConfiguration.Pipeline.ResponseClassifier,
                         Constants.MaxReliabilityRetries);
 
-                    stream = stream.WithNoDispose().WithProgress(progressHandler);
+                    stream = stream.WithProgress(progressHandler);
 
                     /* Decryption handled by caller, so safe to check checksum now.
                      * Buffer response stream and ensure it matches the transactional checksum if any.

diff --git a/sdk/storage/Azure.Storage.Blobs/src/PartitionedDownloader.cs b/sdk/storage/Azure.Storage.Blobs/src/PartitionedDownloader.cs
@@ -6,6 +6,7 @@
 using System.Collections.Generic;
 using System.Globalization;
 using System.IO;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
@@ -144,6 +145,7 @@ public async Task<Response> DownloadToInternal(
             // tracing
             DiagnosticScope scope = _client.ClientConfiguration.ClientDiagnostics.CreateScope(_operationName);
             using DisposableBucket disposables = new DisposableBucket();
+            Queue<Task<Response<BlobDownloadStreamingResult>>> runningTasks = null;
             try
             {
                 scope.Start();
@@ -230,7 +232,6 @@ await HandleOneShotDownload(initialResponse, destination, async, cancellationTok
 #pragma warning disable AZC0110 // DO NOT use await keyword in possibly synchronous scope.
                                 // Rule checker cannot understand this section, but this
                                 // massively reduces code duplication.
-                Queue<Task<Response<BlobDownloadStreamingResult>>> runningTasks = null;
                 int effectiveWorkerCount = async ? _maxWorkerCount : 1;
                 if (effectiveWorkerCount > 1)
                 {
@@ -354,6 +355,17 @@ await CopyToInternal(
             }
             finally
             {
+#pragma warning disable AZC0110
+                if (runningTasks != null)
+                {
+                    async Task DisposeStreamAsync(Task<Response<BlobDownloadStreamingResult>> task)
+                    {
+                        Response<BlobDownloadStreamingResult> response = await task.ConfigureAwait(false);
+                        response.Value.Content.Dispose();
+                    }
+                    await Task.WhenAll(runningTasks.Select(DisposeStreamAsync)).ConfigureAwait(false);
+                }
+#pragma warning restore AZC0110
                 scope.Dispose();
             }
         }