Optimize part streaming

Author: GarrettBeatty

generator/.DevConfigs/9d07dc1e-d82d-4f94-8700-c7b57f872045.json

@@ -0,0 +1,11 @@
+{
+  "services": [
+    {
+      "serviceName": "S3",
+      "type": "patch",
+      "changeLogMessages": [
+        "Optimized multipart download manager to stream responses directly where applicable."
+      ]
+    }
+  ]
+}

sdk/src/Services/S3/Custom/Transfer/Internal/BufferedPartDataHandler.cs

@@ -30,9 +30,26 @@
 namespace Amazon.S3.Transfer.Internal
 {
     /// <summary>
-    /// Buffers downloaded parts in memory using <see cref="ArrayPool{T}"/> and <see cref="IPartBufferManager"/>.
-    /// Implements current streaming behavior for multipart downloads.
+    /// Handles multipart download data with intelligent stream-vs-buffer decision making.
+    /// Optimizes for sequential part arrival by streaming directly to consumer when possible,
+    /// while buffering out-of-order parts into memory using <see cref="ArrayPool{T}"/>.
     /// </summary>
+    /// <remarks>
+    /// <para><strong>Optimization Strategy:</strong></para>
+    /// <list type="bullet">
+    /// <item>Parts arriving in expected order (matching NextExpectedPartNumber) stream directly to consumer</item>
+    /// <item>Out-of-order parts buffer into ArrayPool memory for later sequential consumption</item>
+    /// <item>Best case: All parts in order → zero buffering → pure streaming</item>
+    /// <item>Worst case: All parts out of order → full buffering (original behavior)</item>
+    /// </list>
+    /// 
+    /// 
+    /// <para><strong>Response Ownership:</strong></para>
+    /// <list type="bullet">
+    /// <item>Streaming: StreamingDataSource takes ownership and disposes after reading</item>
+    /// <item>Buffering: Handler disposes response immediately after buffering completes</item>
+    /// </list>
+    /// </remarks>
     internal class BufferedPartDataHandler : IPartDataHandler
     {
         private readonly IPartBufferManager _partBufferManager;
@@ -64,28 +81,163 @@ public Task PrepareAsync(DownloadDiscoveryResult discoveryResult, CancellationTo
         }
 
         /// <inheritdoc/>
+        /// <remarks>
+        /// <para>
+        /// Intelligently chooses between streaming and buffering based on part arrival order:
+        /// </para>
+        /// <list type="bullet">
+        /// <item>If partNumber matches NextExpectedPartNumber: Stream directly (no buffering)</item>
+        /// <item>Otherwise: Buffer into memory for later sequential consumption</item>
+        /// </list>
+        /// <para><strong>Response Ownership:</strong></para>
+        /// <para>
+        /// This method takes ownership of the response and is responsible for disposing it in ALL cases,
+        /// including error scenarios. Callers must NOT dispose the response after calling this method.
+        /// </para>
+        /// </remarks>
         public async Task ProcessPartAsync(
             int partNumber,
             GetObjectResponse response, 
             CancellationToken cancellationToken)
         {
-            Logger.DebugFormat("BufferedPartDataHandler: [Part {0}] Starting to buffer part from response stream - ContentLength={1}",
-                partNumber, response.ContentLength);
+            if (partNumber == _partBufferManager.NextExpectedPartNumber)
+            {
+                await ProcessStreamingPartAsync(partNumber, response, cancellationToken).ConfigureAwait(false);
+            }
+            else
+            {
+                await ProcessBufferedPartAsync(partNumber, response, cancellationToken).ConfigureAwait(false);
+            }
+        }
+        
+        /// <summary>
+        /// Processes a part that arrives in expected order by streaming it directly without buffering.
+        /// Takes ownership of the response and transfers it to the StreamingDataSource.
+        /// </summary>
+        /// <param name="partNumber">The part number being processed.</param>
+        /// <param name="response">The GetObjectResponse containing the part data. Ownership is transferred to StreamingDataSource.</param>
+        /// <param name="cancellationToken">Cancellation token for the operation.</param>
+        /// <remarks>
+        /// This method is called when the part arrives in the expected sequential order, allowing
+        /// for optimal zero-copy streaming directly to the consumer without buffering into memory.
+        /// 
+        /// OWNERSHIP TRANSFER:
+        /// 1. Response is passed to StreamingDataSource constructor (StreamingDataSource takes ownership)
+        /// 2. StreamingDataSource is added to buffer manager (buffer manager takes ownership)
+        /// 3. After successful AddBufferAsync, we null out our reference to mark ownership transfer
+        /// 4. Buffer manager will dispose StreamingDataSource (which disposes response) during cleanup
+        /// 
+        /// ERROR HANDLING:
+        /// - If StreamingDataSource constructor fails: We dispose the response (still our responsibility)
+        /// - If constructor succeeds but AddBufferAsync fails: StreamingDataSource.Dispose() handles the response
+        /// - If AddBufferAsync succeeds: Buffer manager owns everything and will clean up
+        /// </remarks>
+        private async Task ProcessStreamingPartAsync(
+            int partNumber,
+            GetObjectResponse response, 
+            CancellationToken cancellationToken)
+        {
+            Logger.DebugFormat("BufferedPartDataHandler: [Part {0}] Matches NextExpectedPartNumber - streaming directly without buffering",
+                partNumber);
 
-            // Buffer the part from the response stream into memory
-            var buffer = await BufferPartFromResponseAsync(
-                partNumber, 
-                response, 
-                cancellationToken).ConfigureAwait(false);
+            StreamingDataSource streamingDataSource = null;
+            var ownsResponse = true; // Track if we still own the response
+            
+            try
+            {
+                // Create a StreamingDataSource that will stream directly from the response
+                // If successful, StreamingDataSource takes ownership of the response and will dispose it
+                streamingDataSource = new StreamingDataSource(partNumber, response);
+                ownsResponse = false; // Ownership transferred to StreamingDataSource
 
-            Logger.DebugFormat("BufferedPartDataHandler: [Part {0}] Buffered {1} bytes into memory",
-                partNumber, buffer.Length);
+                // Add the streaming data source to the buffer manager
+                // After this succeeds, the buffer manager owns the data source
+                await _partBufferManager.AddBufferAsync(streamingDataSource, cancellationToken).ConfigureAwait(false);
 
-            // Add the buffered part to the buffer manager
-            await _partBufferManager.AddBufferAsync(buffer, cancellationToken).ConfigureAwait(false);
+                // Mark ownership transfer by nulling our reference
+                // If ReleaseBufferSpace() throws, we no longer own the data source, so we won't dispose it
+                streamingDataSource = null;
 
-            Logger.DebugFormat("BufferedPartDataHandler: [Part {0}] Added to buffer manager",
-                partNumber);
+                // Release capacity immediately since we're not holding anything in memory
+                _partBufferManager.ReleaseBufferSpace();
+
+                Logger.DebugFormat("BufferedPartDataHandler: [Part {0}] StreamingDataSource added and capacity released",
+                    partNumber);
+            }
+            catch (Exception ex)
+            {
+                Logger.Error(ex, "BufferedPartDataHandler: [Part {0}] Failed to process streaming part", partNumber);
+                
+                // Dispose response if we still own it (constructor failed before taking ownership)
+                if (ownsResponse)
+                    response?.Dispose();
+                
+                // Dispose StreamingDataSource if we created it but buffer manager doesn't own it yet
+                // If null, the buffer manager owns it and will handle cleanup
+                streamingDataSource?.Dispose();
+                
+                throw;
+            }
+        }
+        
+        /// <summary>
+        /// Processes a part that arrives out of order by buffering it into memory.
+        /// Takes ownership of the response and disposes it after buffering completes.
+        /// </summary>
+        /// <param name="partNumber">The part number being processed.</param>
+        /// <param name="response">The GetObjectResponse containing the part data. This method owns and disposes it.</param>
+        /// <param name="cancellationToken">Cancellation token for the operation.</param>
+        /// <remarks>
+        /// This method is called when the part arrives out of the expected sequential order.
+        /// The part data is buffered into ArrayPool memory for later sequential consumption.
+        /// 
+        /// OWNERSHIP:
+        /// - Response is read and buffered into StreamPartBuffer
+        /// - Response is disposed immediately after buffering (no longer needed)
+        /// - StreamPartBuffer is added to buffer manager (buffer manager takes ownership)
+        /// - Buffer manager will dispose StreamPartBuffer during cleanup
+        /// 
+        /// ERROR HANDLING:
+        /// - Always dispose response in catch block since we own it throughout this method
+        /// - BufferPartFromResponseAsync handles its own cleanup of StreamPartBuffer on error
+        /// </remarks>
+        private async Task ProcessBufferedPartAsync(
+            int partNumber,
+            GetObjectResponse response, 
+            CancellationToken cancellationToken)
+        {
+            Logger.DebugFormat("BufferedPartDataHandler: [Part {0}] Out of order (NextExpected={1}) - buffering to memory",
+                partNumber, _partBufferManager.NextExpectedPartNumber);
+
+            try
+            {
+                // Buffer the part from the response stream into memory
+                var buffer = await BufferPartFromResponseAsync(
+                    partNumber, 
+                    response, 
+                    cancellationToken).ConfigureAwait(false);
+
+                // Response has been fully read and buffered - dispose it now
+                response?.Dispose();
+
+                Logger.DebugFormat("BufferedPartDataHandler: [Part {0}] Buffered {1} bytes into memory",
+                    partNumber, buffer.Length);
+                    
+                // Add the buffered part to the buffer manager
+                await _partBufferManager.AddBufferAsync(buffer, cancellationToken).ConfigureAwait(false);
+
+                Logger.DebugFormat("BufferedPartDataHandler: [Part {0}] Added to buffer manager (capacity will be released after consumption)",
+                    partNumber);
+            }
+            catch (Exception ex)
+            {
+                Logger.Error(ex, "BufferedPartDataHandler: [Part {0}] Failed to process buffered part", partNumber);
+                
+                // We own the response throughout this method, so dispose it on error
+                response?.Dispose();
+                
+                throw;
+            }
         }
 
         /// <inheritdoc/>
@@ -112,6 +264,15 @@ public void Dispose()
             // _partBufferManager is owned by caller, don't dispose
         }
 
+        /// <summary>
+        /// Buffers a part from the GetObjectResponse stream into ArrayPool memory.
+        /// Used when a part arrives out of order and cannot be streamed directly.
+        /// </summary>
+        /// <param name="partNumber">The part number being buffered.</param>
+        /// <param name="response">The GetObjectResponse containing the part data stream.</param>
+        /// <param name="cancellationToken">Cancellation token for the operation.</param>
+        /// <returns>A <see cref="StreamPartBuffer"/> containing the buffered part data.</returns>
+        /// <exception cref="Exception">Thrown when buffering fails. The StreamPartBuffer will be disposed automatically.</exception>
         private async Task<StreamPartBuffer> BufferPartFromResponseAsync(
             int partNumber, 
             GetObjectResponse response, 

sdk/src/Services/S3/Custom/Transfer/Internal/FilePartDataHandler.cs

@@ -69,26 +69,42 @@ public Task PrepareAsync(DownloadDiscoveryResult discoveryResult, CancellationTo
         }
 
         /// <inheritdoc/>
+        /// <remarks>
+        /// <para><strong>Response Ownership:</strong></para>
+        /// <para>
+        /// This method takes ownership of the response and is responsible for disposing it in ALL cases,
+        /// including error scenarios. Callers must NOT dispose the response after calling this method.
+        /// </para>
+        /// </remarks>
         public async Task ProcessPartAsync(
             int partNumber,
             GetObjectResponse response,
             CancellationToken cancellationToken)
         {
-            Logger.DebugFormat("FilePartDataHandler: [Part {0}] Starting to process part - ContentLength={1}",
-                partNumber, response.ContentLength);
+            try
+            {
+                Logger.DebugFormat("FilePartDataHandler: [Part {0}] Starting to process part - ContentLength={1}",
+                    partNumber, response.ContentLength);
 
-            // Calculate offset for this part based on ContentRange or part number
-            long offset = GetPartOffset(response, partNumber);
+                // Calculate offset for this part based on ContentRange or part number
+                long offset = GetPartOffset(response, partNumber);
 
-            Logger.DebugFormat("FilePartDataHandler: [Part {0}] Calculated file offset={1}",
-                partNumber, offset);
+                Logger.DebugFormat("FilePartDataHandler: [Part {0}] Calculated file offset={1}",
+                    partNumber, offset);
 
-            // Write part data to file at the calculated offset
-            await WritePartToFileAsync(offset, response, cancellationToken)
-                .ConfigureAwait(false);
+                // Write part data to file at the calculated offset
+                await WritePartToFileAsync(offset, response, cancellationToken)
+                    .ConfigureAwait(false);
 
-            Logger.DebugFormat("FilePartDataHandler: [Part {0}] File write completed successfully",
-                partNumber);
+                Logger.DebugFormat("FilePartDataHandler: [Part {0}] File write completed successfully",
+                    partNumber);
+            }
+            finally
+            {
+                // Always dispose response after writing to disk (success or failure)
+                // This releases the HTTP connection back to the pool
+                response?.Dispose();
+            }
         }
 
         /// <inheritdoc/>

sdk/src/Services/S3/Custom/Transfer/Internal/IPartBufferManager.cs

@@ -46,6 +46,14 @@ internal interface IPartBufferManager : IDisposable
         /// <returns>A task that completes when the buffer has been added and signaling is complete.</returns>
         Task AddBufferAsync(StreamPartBuffer buffer, CancellationToken cancellationToken);
 
+        /// <summary>
+        /// Adds a part data source (streaming or buffered) and signals readers when next expected part arrives.
+        /// </summary>
+        /// <param name="dataSource">The part data source to add (can be StreamingDataSource or BufferedDataSource).</param>
+        /// <param name="cancellationToken">A token to cancel the operation.</param>
+        /// <returns>A task that completes when the data source has been added and signaling is complete.</returns>
+        Task AddBufferAsync(IPartDataSource dataSource, CancellationToken cancellationToken);
+        
         /// <summary>
         /// Reads data from the buffer manager. Automatically handles sequential part consumption
         /// and reads across part boundaries to fill the buffer when possible, matching standard Stream.Read() behavior.

sdk/src/Services/S3/Custom/Transfer/Internal/MultipartDownloadManager.cs

@@ -387,25 +387,26 @@ private async Task CreateDownloadTaskAsync(int partNumber, long objectSize, Even
                         partNumber, _httpConcurrencySlots.CurrentCount, _config.ConcurrentServiceRequests);
                 }
 
-                Logger.DebugFormat("MultipartDownloadManager: [Part {0}] Starting buffering", partNumber);
+                Logger.DebugFormat("MultipartDownloadManager: [Part {0}] Processing part (handler will decide: stream or buffer)", partNumber);
 
                 // Delegate data handling to the handler
+                // IMPORTANT: Handler takes ownership of response and is responsible for disposing it in ALL cases:
+                // - If streaming: StreamingDataSource takes ownership and disposes when consumer finishes reading
+                // - If buffering: Handler disposes immediately after copying data to buffer
+                // - On error: Handler disposes in its catch block before rethrowing
                 await _dataHandler.ProcessPartAsync(partNumber, response, cancellationToken).ConfigureAwait(false);
 
-                Logger.DebugFormat("MultipartDownloadManager: [Part {0}] Buffering completed successfully", partNumber);
+                Logger.DebugFormat("MultipartDownloadManager: [Part {0}] Processing completed successfully", partNumber);
             }
             catch (Exception ex)
             {
                 Logger.Error(ex, "MultipartDownloadManager: [Part {0}] Download failed", partNumber);
+                
                 // Release capacity on failure
+                // Note: Handler is responsible for disposing response in all cases, including errors
                 _dataHandler.ReleaseCapacity();
                 throw;
             }
-            finally
-            {
-                // Always dispose the response since we never transfer ownership
-                response?.Dispose();
-            }
         }