prepare for cancellation handling

Author: ktoso

include/swift/Runtime/Concurrency.h

@@ -307,7 +307,8 @@ bool swift_taskGroup_isDiscardingResults(TaskGroup *group);
 /// \code
 /// func swift_taskGroup_waitAll(
 ///     waitingTask: Builtin.NativeObject, // current task
-///     group: Builtin.RawPointer
+///     group: Builtin.RawPointer,
+///     childFailureCancelsGroup: Bool
 /// ) async throws
 /// \endcode
   SWIFT_EXPORT_FROM(swift_Concurrency)
@@ -316,6 +317,7 @@ bool swift_taskGroup_isDiscardingResults(TaskGroup *group);
       OpaqueValue *resultPointer,
       SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
       TaskGroup *group,
+      bool childFailureCancelsGroup,
       ThrowingTaskFutureWaitContinuationFunction *resumeFn,
       AsyncContext *callContext);
 

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -326,10 +326,11 @@ OVERRIDE_TASK_GROUP(taskGroup_waitAll, void,
                     (OpaqueValue *resultPointer,
                         SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
                         TaskGroup *_group,
+                        bool childFailureCancelsGroup,
                         ThrowingTaskFutureWaitContinuationFunction *resumeFn,
                         AsyncContext *callContext),
-                    (resultPointer, callerContext, _group, resumeFn,
-                        callContext))
+                    (resultPointer, callerContext, _group, childFailureCancelsGroup,
+                        resumeFn, callContext))
 
 OVERRIDE_TASK_LOCAL(task_reportIllegalTaskLocalBindingWithinWithTaskGroup, void,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,

stdlib/public/Concurrency/TaskGroup.cpp

@@ -54,7 +54,7 @@ using namespace swift;
 /*************************** TASK GROUP ***************************************/
 /******************************************************************************/
 
-#if 1
+#if 0
 #define SWIFT_TASK_GROUP_DEBUG_LOG(group, fmt, ...) \
 fprintf(stderr, "[%#lx] [%s:%d](%s) group(%p%s) " fmt "\n",             \
       (unsigned long)Thread::current().platformThreadId(),              \
@@ -956,6 +956,7 @@ __attribute__((noinline))
 SWIFT_CC(swiftasync) static void workaround_function_swift_taskGroup_waitAllImpl(
     OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
     TaskGroup *_group,
+    bool childFailureCancelsGroup,
     ThrowingTaskFutureWaitContinuationFunction resumeFunction,
     AsyncContext *callContext) {
   // Make sure we don't eliminate calls to this function.
@@ -1187,6 +1188,7 @@ SWIFT_CC(swiftasync)
 static void swift_taskGroup_waitAllImpl(
     OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
     TaskGroup *_group,
+    bool childFailureCancelsGroup,
     ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
     AsyncContext *rawContext) {
   auto waitingTask = swift_task_getCurrent();
@@ -1213,7 +1215,7 @@ static void swift_taskGroup_waitAllImpl(
       // there were pending tasks so it will be woken up eventually.
 #ifdef __ARM_ARCH_7K__
       return workaround_function_swift_taskGroup_waitAllImpl(
-         resultPointer, callerContext, _group, resumeFunction, rawContext);
+         resultPointer, callerContext, _group, childFailureCancelsGroup, resumeFunction, rawContext);
 #else /* __ARM_ARCH_7K__ */
       return;
 #endif /* __ARM_ARCH_7K__ */

stdlib/public/Concurrency/TaskGroup.swift

@@ -110,7 +110,7 @@ public func withTaskGroup<ChildTaskResult, GroupResult>(
   // Run the withTaskGroup body.
   let result = await body(&group)
 
-  let _: ChildTaskResult? = try? await _taskGroupWaitAll(group: _group) // try!-safe, cannot throw since this is a non throwing group
+  let _: ChildTaskResult? = try? await _taskGroupWaitAll(group: _group, childFailureCancelsGroup: discardResults) // try!-safe, cannot throw since this is a non throwing group
   return result
 }
 
@@ -198,14 +198,14 @@ public func withThrowingTaskGroup<ChildTaskResult, GroupResult>(
     // Run the withTaskGroup body.
     let result = try await body(&group)
 
-    _ = try? await group.awaitAllRemainingTasks()
+    await group.awaitAllRemainingTasks()
     Builtin.destroyTaskGroup(_group)
 
     return result
   } catch {
     group.cancelAll()
 
-    _ = try? await group.awaitAllRemainingTasks() // discard errors
+    await group.awaitAllRemainingTasks()
     Builtin.destroyTaskGroup(_group)
 
     throw error
@@ -231,28 +231,24 @@ public func withThrowingTaskGroup<ChildTaskResult, GroupResult>(
 
   let _group = Builtin.createTaskGroupWithFlags(flags, ChildTaskResult.self)
   var group = ThrowingTaskGroup<ChildTaskResult, Error>(group: _group)
+  defer { Builtin.destroyTaskGroup(_group) }
 
+  let result: GroupResult
   do {
     // Run the withTaskGroup body.
-    let result = try await body(&group)
-
-    try await group.awaitAllRemainingTasks()
-    Builtin.destroyTaskGroup(_group)
-
-    return result
+    result = try await body(&group)
   } catch {
     group.cancelAll()
 
-    do {
-      try await group.awaitAllRemainingTasks()
-      Builtin.destroyTaskGroup(_group)
-    } catch {
-      Builtin.destroyTaskGroup(_group)
-      throw error
-    }
+    await group.awaitAllRemainingTasks()
 
     throw error
   }
+
+  // FIXME: if one of them throws, cancel the group
+  try await group.awaitAllRemainingTasksThrowing(childFailureCancelsGroup: true)
+
+  return result
 }
 
 /// A group that contains dynamically created child tasks.
@@ -514,7 +510,7 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
     /// implementation.
     if #available(SwiftStdlib 5.8, *) {
       if isDiscardingResults {
-        let _: ChildTaskResult? = try! await _taskGroupWaitAll(group: _group) // try!-safe, cannot throw, not throwing group
+        let _: ChildTaskResult? = try! await _taskGroupWaitAll(group: _group, childFailureCancelsGroup: isDiscardingResults) // try!-safe, cannot throw, not throwing group
         return
       }
     }
@@ -641,13 +637,22 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
 
   /// Await all the remaining tasks on this group.
   @usableFromInline
-  internal mutating func awaitAllRemainingTasks() async throws {
+  @available(*, deprecated, message: "Use `awaitAllRemainingTasksThrowing`, since 5.8 with discardResults draining may throw")
+  internal mutating func awaitAllRemainingTasks() async {
+    // We discard the error because in old code, which may have inlined this `awaitAllRemainingTasks`
+    // method, draining was never going to throw
+    _ = try? await awaitAllRemainingTasksThrowing(childFailureCancelsGroup: false)
+  }
+
+  /// Await all the remaining tasks on this group.
+  @usableFromInline
+  internal mutating func awaitAllRemainingTasksThrowing(childFailureCancelsGroup: Bool) async throws {
     /// Since 5.8, we implement "wait for all pending tasks to complete"
     /// in the runtime, in order to be able to handle the discard-results
     /// implementation.
     if #available(SwiftStdlib 5.8, *) {
       if isDiscardingResults {
-        let _: ChildTaskResult? = try await _taskGroupWaitAll(group: _group) // if any of the tasks throws, this will "rethrow" here
+        let _: ChildTaskResult? = try await _taskGroupWaitAll(group: _group, childFailureCancelsGroup: childFailureCancelsGroup) // if any of the tasks throws, this will "rethrow" here
         return
       }
     }
@@ -674,7 +679,7 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
   /// - Throws: only during
   @_alwaysEmitIntoClient
   public mutating func waitForAll() async throws {
-    try await self.awaitAllRemainingTasks()
+    try await self.awaitAllRemainingTasksThrowing(childFailureCancelsGroup: false)
   }
 
 #if !SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY
@@ -1207,7 +1212,7 @@ func _taskHasTaskGroupStatusRecord() -> Bool
 @usableFromInline
 @discardableResult
 @_silgen_name("swift_taskGroup_waitAll")
-func _taskGroupWaitAll<T>(group: Builtin.RawPointer) async throws -> T?
+func _taskGroupWaitAll<T>(group: Builtin.RawPointer, childFailureCancelsGroup: Bool) async throws -> T?
 
 @available(SwiftStdlib 5.8, *)
 @_silgen_name("swift_taskGroup_isDiscardingResults")

stdlib/public/Concurrency/TaskPrivate.h

@@ -137,7 +137,10 @@ namespace {
 ///   func _taskGroupWaitNext<T>(group: Builtin.RawPointer) async throws -> T?
 ///
 ///   @_silgen_name("swift_taskGroup_waitAll")
-///   func _taskGroupWaitAll<T>(group: Builtin.RawPointer) async throws -> T?
+///   func _taskGroupWaitAll<T>(
+///     group: Builtin.RawPointer,
+///     childFailureCancelsGroup: Bool
+///   ) async throws -> T?
 ///
 class TaskFutureWaitAsyncContext : public AsyncContext {
 public:

test/Concurrency/Runtime/async_taskgroup_throw_rethrow.swift

@@ -93,14 +93,17 @@ func test_taskGroup_discardResults_automaticallyRethrowsOnlyFirst() async {
     let got = try await withThrowingTaskGroup(of: Int.self, returning: Int.self,
                                               discardResults: true) { group in
       group.addTask { await echo(1) }
-      group.addTask { throw Boom(id: "first") }
+      group.addTask { throw Boom(id: "first-a") }
+      group.addTask { throw Boom(id: "first-b") }
       // add a throwing task, but don't consume it explicitly
       // since we're in discard results mode, all will be awaited and the first error it thrown
 
       do {
         try await group.waitForAll()
       } catch {
-        // CHECK: caught: Boom(id: "first")
+        // There's no guarantee about which of the `first-...` tasks will complete first,
+        // however, they all will be consumed when we have returned from the `waitForAll`.
+        // CHECK: caught: Boom(id: "first
         print("caught: \(error)")
       }
 

unittests/runtime/CompatibilityOverrideConcurrency.cpp

@@ -216,7 +216,7 @@ TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_taskGroup_cancelAll) {
 }
 
 TEST_F(CompatibilityOverrideConcurrencyTest, test_swift_taskGroup_waitAll) {
-  swift_taskGroup_waitAll(nullptr, nullptr, nullptr, nullptr,
+  swift_taskGroup_waitAll(nullptr, nullptr, nullptr, false, nullptr,
                           nullptr);
 }