Fix #63147 waitForAll impl in ThrowingTaskGroup, used to not wait enough

rdar://104507347

Author: ktoso

stdlib/public/Concurrency/TaskGroup.swift

@@ -616,11 +616,25 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
   }
 
   /// Wait for all of the group's remaining tasks to complete.
-  ///
-  /// - Throws: only during
   @_alwaysEmitIntoClient
   public mutating func waitForAll() async throws {
-    await self.awaitAllRemainingTasks()
+    var firstError: Error? = nil
+
+    // Make sure we loop until all child tasks have completed
+    while !isEmpty {
+      do {
+        while let _ = try await next() {}
+      } catch {
+        // Upon error throws, capture the first one
+        if firstError == nil {
+          firstError = error
+        }
+      }
+    }
+
+    if let firstError {
+      throw firstError
+    }
   }
 
 #if !SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY

test/Concurrency/Runtime/async_taskgroup_throw_rethrow.swift

@@ -68,6 +68,33 @@ func test_taskGroup_noThrow_ifNotAwaitedThrowingTask() async {
   print("Expected no error to be thrown, got: \(got)") // CHECK: Expected no error to be thrown, got: 1
 }
 
+func test_taskGroup_throw_rethrows_waitForAll() async {
+  print("==== \(#function) ------") // CHECK-LABEL: test_taskGroup_throw_rethrows_waitForAll
+  do {
+    _ = try await withThrowingTaskGroup(of: Int.self) { group in
+      group.addTask {
+        throw CancellationError()
+      }
+      group.addTask {
+        1
+      }
+
+      do {
+        try await group.waitForAll()
+      } catch {
+        print("waitAll rethrown: ", error)
+        // CHECK: waitAll rethrown: CancellationError()
+        print("isEmpty: ", group.isEmpty)
+        // CHECK: isEmpty: true
+        throw error
+      }
+    }
+  } catch {
+    print("rethrown: ", error)
+    // CHECK: rethrown: CancellationError()
+  }
+}
+
 func test_discardingTaskGroup_automaticallyRethrows() async {
   print("==== \(#function) ------") // CHECK-LABEL: test_discardingTaskGroup_automaticallyRethrows
   do {
@@ -182,6 +209,7 @@ func test_discardingTaskGroup_automaticallyRethrows_first_withThrowingBodySecond
   static func main() async {
     await test_taskGroup_throws_rethrows()
     await test_taskGroup_noThrow_ifNotAwaitedThrowingTask()
+    await test_taskGroup_throw_rethrows_waitForAll()
     await test_discardingTaskGroup_automaticallyRethrows()
     await test_discardingTaskGroup_automaticallyRethrowsOnlyFirst()
     await test_discardingTaskGroup_automaticallyRethrows_first_withThrowingBodyFirst()