[Concurrency] Cancelled group should only spawn already cancelled tasks (swiftlang#38073)

Author: ktoso

include/swift/ABI/TaskGroup.h

@@ -39,6 +39,9 @@ class alignas(Alignment_TaskGroup) TaskGroup {
 
   /// Upon a future task's completion, offer it to the task group it belongs to.
   void offer(AsyncTask *completed, AsyncContext *context);
+
+  /// Checks the cancellation status of the group.
+  bool isCancelled();
 };
 
 } // end namespace swift

stdlib/public/Concurrency/Task.cpp

@@ -557,7 +557,8 @@ static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
     // In a task group we would not have allowed the `add` to create a child anymore,
     // however better safe than sorry and `async let` are not expressed as task groups,
     // so they may have been spawned in any case still.
-    if (swift_task_isCancelled(parent))
+    if (swift_task_isCancelled(parent) ||
+        (group && group->isCancelled()))
       swift_task_cancel(task);
 
     // Initialize task locals with a link to the parent task.

stdlib/public/Concurrency/TaskGroup.cpp

@@ -508,6 +508,10 @@ void TaskGroup::offer(AsyncTask *completedTask, AsyncContext *context) {
   asImpl(this)->offer(completedTask, context);
 }
 
+bool TaskGroup::isCancelled() {
+  return asImpl(this)->isCancelled();
+}
+
 static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
                                 PollResult result) {
   /// Fill in the result value

stdlib/public/Concurrency/TaskGroup.swift

@@ -658,12 +658,12 @@ public struct ThrowingTaskGroup<ChildTaskResult, Failure: Error> {
     _taskGroupIsEmpty(_group)
   }
 
-  /// Cancel all of the remaining tasks in the group.
-  ///
-  /// After canceling a group, adding a new task to it always fails.
+  /// Cancel all the remaining, and future, tasks in the group.
   ///
-  /// Any results, including errors thrown by tasks affected by this
-  /// cancellation, are silently discarded.
+  /// A cancelled group will not will create new tasks when the `asyncUnlessCancelled`,
+  /// function is used. It will, however, continue to create tasks when the plain `async`
+  /// function is used. Such tasks will be created yet immediately cancelled, allowing
+  /// the tasks to perform some short-cut implementation, if they are responsive to cancellation.
   ///
   /// This function may be called even from within child (or any other) tasks,
   /// and causes the group to be canceled.

test/Concurrency/Runtime/async_taskgroup_cancelAll_only_specific_group.swift

@@ -22,7 +22,7 @@ func test_taskGroup_cancelAll_onlySpecificGroup() async {
   async let g1: Int = withTaskGroup(of: Int.self) { group in
 
     for i in 1...5 {
-      group.spawn {
+      group.async {
         await Task.sleep(1_000_000_000)
         let c = Task.isCancelled
         print("add: \(i) (cancelled: \(c))")
@@ -31,7 +31,7 @@ func test_taskGroup_cancelAll_onlySpecificGroup() async {
     }
 
     var sum = 0
-    while let got = try! await group.next() {
+    while let got = await group.next() {
       print("next: \(got)")
       sum += got
     }
@@ -45,9 +45,9 @@ func test_taskGroup_cancelAll_onlySpecificGroup() async {
   }
 
   // The cancellation os g2 should have no impact on g1
-  let g2: Int = try! await withTaskGroup(of: Int.self) { group in
+  let g2: Int = await withTaskGroup(of: Int.self) { group in
     for i in 1...3 {
-      group.spawn {
+      group.async {
         await Task.sleep(1_000_000_000)
         let c = Task.isCancelled
         print("g1 task \(i) (cancelled: \(c))")
@@ -65,8 +65,8 @@ func test_taskGroup_cancelAll_onlySpecificGroup() async {
     return 0
   }
 
-  let result1 = try! await g1
-  let result2 = try! await g2
+  let result1 = await g1
+  let result2 = g2
 
   // CHECK: g2 task cancelled: false
   // CHECK: g2 group cancelled: true

test/Concurrency/Runtime/async_taskgroup_cancel_then_spawn.swift

@@ -37,13 +37,16 @@ func test_taskGroup_cancel_then_add() async {
     let none = await group.next()
     print("next second: \(none)") // CHECK: next second: nil
 
-    group.spawn { 3 }
-    print("added third, unconditionally") // CHECK: added third, unconditionally
-    print("group isCancelled: \(group.isCancelled)") // CHECK: group isCancelled: true
-
+    group.spawn {
+      print("child task isCancelled: \(Task.isCancelled)") // CHECK: child task isCancelled: true
+      return 3
+    }
     let three = await group.next()!
     print("next third: \(three)") // CHECK: next third: 3
 
+    print("added third, unconditionally") // CHECK: added third, unconditionally
+    print("group isCancelled: \(group.isCancelled)") // CHECK: group isCancelled: true
+
     return one + (none ?? 0)
   }