queue fixes

Author: GarrettBeatty

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

@@ -414,7 +414,9 @@ private async Task StartBackgroundDownloadsAsync(DownloadResult downloadResult,
                 _logger.DebugFormat("MultipartDownloadManager: Background task waiting for {0} download tasks", expectedTaskCount);
 
                 // Wait for all downloads to complete (fails fast on first exception)
-                await TaskHelpers.WhenAllOrFirstExceptionAsync(downloadTasks, internalCts.Token).ConfigureAwait(false);
+                // Use fault-priority variant to ensure original exceptions propagate instead of
+                // OperationCanceledException when cancellation occurs due to a task failure
+                await TaskHelpers.WhenAllOrFirstExceptionWithFaultPriorityAsync(downloadTasks, internalCts.Token).ConfigureAwait(false);
 
                 _logger.DebugFormat("MultipartDownloadManager: All download tasks completed successfully");
 
@@ -451,13 +453,21 @@ private async Task CreateDownloadTasksAsync(DownloadResult downloadResult, Event
             // Pre-acquire capacity in sequential order to prevent race condition deadlock
             // This ensures Part 2 gets capacity before Part 3, etc., preventing out-of-order
             // parts from consuming all buffer slots and blocking the next expected part
-            for (int partNum = 2; partNum <= downloadResult.TotalParts; partNum++)
+            for (int partNum = 2; partNum <= downloadResult.TotalParts && !internalCts.IsCancellationRequested; partNum++)
             {
                 _logger.DebugFormat("MultipartDownloadManager: [Part {0}] Waiting for buffer space", partNum);
 
                 // Acquire capacity sequentially - guarantees Part 2 before Part 3, etc.
                 await _dataHandler.WaitForCapacityAsync(internalCts.Token).ConfigureAwait(false);
 
+                // Check cancellation after acquiring capacity - a task may have failed while waiting
+                if (internalCts.IsCancellationRequested)
+                {
+                    _logger.InfoFormat("MultipartDownloadManager: [Part {0}] Stopping early - cancellation requested after capacity acquired", partNum);
+                    _dataHandler.ReleaseCapacity();
+                    break;
+                }
+                
                 _logger.DebugFormat("MultipartDownloadManager: [Part {0}] Buffer space acquired", partNum);
 
                 _logger.DebugFormat("MultipartDownloadManager: [Part {0}] Waiting for HTTP concurrency slot (Available: {1}/{2})",
@@ -466,22 +476,58 @@ private async Task CreateDownloadTasksAsync(DownloadResult downloadResult, Event
                 // Acquire HTTP slot in the loop before creating task
                 // Loop will block here if all slots are in use
                 await _httpConcurrencySlots.WaitAsync(internalCts.Token).ConfigureAwait(false);
+                
+                // Check cancellation after acquiring HTTP slot - a task may have failed while waiting
+                if (internalCts.IsCancellationRequested)
+                {
+                    _logger.InfoFormat("MultipartDownloadManager: [Part {0}] Stopping early - cancellation requested after HTTP slot acquired", partNum);
+                    _httpConcurrencySlots.Release();
+                    _dataHandler.ReleaseCapacity();
+                    break;
+                }
 
                 _logger.DebugFormat("MultipartDownloadManager: [Part {0}] HTTP concurrency slot acquired", partNum);
 
                 try
                 {
                     var task = CreateDownloadTaskAsync(partNum, downloadResult.ObjectSize, wrappedCallback, internalCts.Token);
+                    
+                    // Add failure detection to immediately cancel internal token on first error
+                    // This prevents the for loop from queuing additional parts after a failure
+                    _ = task.ContinueWith(t =>
+                    {
+                        if (t.IsFaulted && !internalCts.IsCancellationRequested)
+                        {
+                            // Then cancel to stop queuing more parts
+                            // Note: The original exception will be propagated by WhenAllOrFirstExceptionAsync
+                            try
+                            {
+                                internalCts.Cancel();
+                                _logger.DebugFormat("MultipartDownloadManager: [Part {0}] Cancelled internal token to stop queuing", partNum);
+                            }
+                            catch (ObjectDisposedException)
+                            {
+                                _logger.DebugFormat("MultipartDownloadManager: [Part {0}] CancellationTokenSource already disposed during cancellation", partNum);
+                            }
+                        }
+                    }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
+                    
                     downloadTasks.Add(task);
                 }
                 catch (Exception ex)
                 {
                     // If task creation fails, release the HTTP slot we just acquired
                     _httpConcurrencySlots.Release();
+                    _dataHandler.ReleaseCapacity();
                     _logger.DebugFormat("MultipartDownloadManager: [Part {0}] HTTP concurrency slot released due to task creation failure: {1}", partNum, ex);
                     throw;
                 }
             }
+            
+            if (internalCts.IsCancellationRequested && downloadTasks.Count < downloadResult.TotalParts - 1)
+            {
+                _logger.InfoFormat("MultipartDownloadManager: Stopped queuing early at {0} parts due to cancellation", downloadTasks.Count);
+            }
         }
 
         /// <summary>

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

@@ -27,6 +27,53 @@ namespace Amazon.S3.Transfer.Internal
     /// </summary>
     internal static class TaskHelpers
     {
+        /// <summary>
+        /// Waits for all tasks to complete or till any task fails or is canceled.
+        /// Prioritizes returning already-completed (especially faulted) tasks before checking cancellation.
+        /// This ensures the original exception is propagated rather than OperationCanceledException when
+        /// cancellation occurs due to a task failure.
+        /// </summary>
+        /// <param name="pendingTasks">List of tasks to wait for completion</param>
+        /// <param name="cancellationToken">Cancellation token to observe</param>
+        /// <returns>A task that represents the completion of all tasks or the first exception</returns>
+        internal static async Task WhenAllOrFirstExceptionWithFaultPriorityAsync(List<Task> pendingTasks, CancellationToken cancellationToken)
+        {
+            int processed = 0;
+            int total = pendingTasks.Count;
+            
+            Logger.GetLogger(typeof(TaskHelpers)).DebugFormat("TaskHelpers.WhenAllOrFirstExceptionWithFaultPriorityAsync: Starting with TotalTasks={0}", total);
+            
+            while (processed < total)
+            {
+                // First check if any task has already completed (synchronously)
+                // This ensures faulted tasks are processed before cancellation is checked,
+                // so the original exception takes precedence over OperationCanceledException
+                var completedTask = pendingTasks.FirstOrDefault(t => t.IsCompleted);
+                
+                if (completedTask == null)
+                {
+                    // No task completed yet - check cancellation before blocking
+                    cancellationToken.ThrowIfCancellationRequested();
+                    
+                    // Now wait for one to complete
+                    completedTask = await Task.WhenAny(pendingTasks)
+                        .ConfigureAwait(continueOnCapturedContext: false);
+                }
+                
+                // Process the completed task - this will throw if faulted
+                await completedTask
+                    .ConfigureAwait(continueOnCapturedContext: false);
+                
+                pendingTasks.Remove(completedTask);
+                processed++;
+                
+                Logger.GetLogger(typeof(TaskHelpers)).DebugFormat("TaskHelpers.WhenAllOrFirstExceptionWithFaultPriorityAsync: Task completed (Processed={0}/{1}, Remaining={2})",
+                    processed, total, pendingTasks.Count);
+            }
+            
+            Logger.GetLogger(typeof(TaskHelpers)).DebugFormat("TaskHelpers.WhenAllOrFirstExceptionWithFaultPriorityAsync: All tasks completed (Total={0})", total);
+        }
+
         /// <summary>
         /// Waits for all tasks to complete or till any task fails or is canceled.
         /// </summary>