Avoid OOM in Task.AddExceptionsForCompletedTask during stress runs (#136268)

src/libraries/System.Private.CoreLib/src/System/Exception.cs

@@ -233,13 +233,19 @@ private string GetStackTrace()
         }
 
         [StackTraceHidden]
-        internal void SetCurrentStackTrace()
+        internal void SetCurrentStackTrace(Exception? source = null)
         {
             if (!CanSetRemoteStackTrace())
             {
                 return; // early-exit
             }
 
+            if (source?._remoteStackTraceString is string existingTrace)
+            {
+                _remoteStackTraceString = existingTrace;
+                return;
+            }
+
             // Store the current stack trace into the "remote" stack trace, which was originally introduced to support
             // remoting of exceptions cross app-domain boundaries, and is thus concatenated into Exception.StackTrace
             // when it's retrieved.

src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.MultiThreading.cs

@@ -249,7 +249,8 @@ private static bool WaitAllCore(ReadOnlySpan<Task> tasks, int millisecondsTimeou
                 if (!exceptionSeen) cancellationToken.ThrowIfCancellationRequested();
 
                 // Now gather up and throw all of the exceptions.
-                foreach (Task task in tasks) AddExceptionsForCompletedTask(ref exceptions, task);
+                TaskCanceledException? sharedCanceledException = null;
+                foreach (Task task in tasks) AddExceptionsForCompletedTask(ref exceptions, task, ref sharedCanceledException);
                 Debug.Assert(exceptions != null, "Should have seen at least one exception");
                 ThrowHelper.ThrowAggregateException(exceptions);
             }

src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.SingleThreading.cs

@@ -95,9 +95,10 @@ private static bool WaitAllCore(ReadOnlySpan<Task> tasks, int millisecondsTimeou
             }
 
             List<Exception>? exceptions = null;
+            TaskCanceledException? sharedCanceledException = null;
             foreach (Task task in tasks)
             {
-                AddExceptionsForCompletedTask(ref exceptions, task);
+                AddExceptionsForCompletedTask(ref exceptions, task, ref sharedCanceledException);
             }
             Debug.Assert(exceptions is not null);
             ThrowHelper.ThrowAggregateException(exceptions);

src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs

@@ -4994,9 +4994,19 @@ private static void AddToList<T>(T item, ref List<T>? list, int initSize)
         /// </summary>
         internal static void AddExceptionsForCompletedTask(ref List<Exception>? exceptions, Task t)
         {
-            AggregateException? ex = t.GetExceptions(true);
-            if (ex != null)
+            TaskCanceledException? sharedCanceledException = null;
+            AddExceptionsForCompletedTask(ref exceptions, t, ref sharedCanceledException);
+        }
+
+        internal static void AddExceptionsForCompletedTask(
+            ref List<Exception>? exceptions,
+            Task t,
+            ref TaskCanceledException? sharedCanceledException)
+        {
+            if (t.ExceptionRecorded)
             {
+                AggregateException ex = t.GetExceptions(true)!;
+
                 // make sure the task's exception observed status is set appropriately
                 // it's possible that WaitAll was called by the parent of an attached child,
                 // this will make sure it won't throw again in the implicit wait
@@ -5005,6 +5015,30 @@ internal static void AddExceptionsForCompletedTask(ref List<Exception>? exceptio
                 exceptions ??= new List<Exception>(ex.InnerExceptionCount);
                 exceptions.AddRange(ex.InternalInnerExceptions);
             }
+            else if (t.IsCanceled)
+            {
+                t.UpdateExceptionObservedStatus();
+
+                exceptions ??= new List<Exception>();
+
+                // Don't go through GetExceptions here: for a canceled task it allocates a
+                // TaskCanceledException, captures its stack trace, and wraps it in an
+                // AggregateException just to unwrap it again.
+                // To avoid OOM during mass cancellations (e.g. in WaitAll), capture the stack trace
+                // once and share it across all synthetic TaskCanceledException instances.
+                TaskCanceledException tce = new TaskCanceledException(t);
+                if (sharedCanceledException == null)
+                {
+                    tce.SetCurrentStackTrace();
+                    sharedCanceledException = tce;
+                }
+                else
+                {
+                    tce.SetCurrentStackTrace(sharedCanceledException);
+                }
+
+                exceptions.Add(tce);
+            }
         }
 
         /// <summary>

src/libraries/System.Runtime/tests/System.Threading.Tasks.Tests/Task/TaskWaitAllAnyTest.cs

@@ -1293,5 +1293,23 @@ public static void TaskWaitAll_Enumerable_AllComplete()
             Task.WaitAll((List<Task>)[Task.Delay(1)]);
             Task.WaitAll((IEnumerable<Task>)[Task.Delay(1), Task.CompletedTask, Task.Delay(1), Task.CompletedTask, Task.Delay(1)]);
         }
+
+        [Fact]
+        public static void TaskWaitAll_CanceledTasks_PreservesStackTrace()
+        {
+            using var cts = new CancellationTokenSource();
+            cts.Cancel();
+            Task t1 = Task.FromCanceled(cts.Token);
+            Task t2 = Task.FromCanceled(cts.Token);
+
+            AggregateException ex = Assert.Throws<AggregateException>(() => Task.WaitAll(t1, t2));
+            Assert.Equal(2, ex.InnerExceptions.Count);
+            foreach (Exception inner in ex.InnerExceptions)
+            {
+                var tce = Assert.IsAssignableFrom<TaskCanceledException>(inner);
+                Assert.NotNull(tce.StackTrace);
+                Assert.NotEmpty(tce.StackTrace);
+            }
+        }
     }
 }