[Concurrency] Parent task cancellation must cancel task group itself.

include/swift/ABI/TaskGroup.h

@@ -44,6 +44,12 @@ class alignas(Alignment_TaskGroup) TaskGroup {
   /// Checks the cancellation status of the group.
   bool isCancelled();
 
+  /// Only mark the task group as cancelled, without performing the follow-up
+  /// work of cancelling all the child tasks.
+  ///
+  /// Returns true if the group was already cancelled before this call.
+  bool statusCancel();
+
   // Add a child task to the task group. Always called while holding the
   // status record lock of the task group's owning task.
   void addChildTask(AsyncTask *task);

stdlib/public/Concurrency/TaskGroup.cpp

@@ -1166,6 +1166,10 @@ bool TaskGroup::isCancelled() {
   return asBaseImpl(this)->isCancelled();
 }
 
+bool TaskGroup::statusCancel() {
+  return asBaseImpl(this)->statusCancel();
+}
+
 // =============================================================================
 // ==== offer ------------------------------------------------------------------
 
@@ -2116,8 +2120,8 @@ bool TaskGroupBase::cancelAll(AsyncTask *owningTask) {
 
   // Cancel all the child tasks.  TaskGroup is not a Sendable type,
   // so cancelAll() can only be called from the owning task.  This
-  // satisfies the precondition on cancelAllChildren_unlocked().
-  _swift_taskGroup_cancelAllChildren_unlocked(asAbstract(this), owningTask);
+  // satisfies the precondition on cancel_unlocked().
+  _swift_taskGroup_cancel_unlocked(asAbstract(this), owningTask);
 
   return true;
 }
@@ -2126,8 +2130,8 @@ SWIFT_CC(swift)
 static void swift_task_cancel_group_child_tasksImpl(TaskGroup *group) {
   // TaskGroup is not a Sendable type, and so this operation (which is not
   // currently exposed in the API) can only be called from the owning
-  // task.  This satisfies the precondition on cancelAllChildren_unlocked().
-  _swift_taskGroup_cancelAllChildren_unlocked(group, swift_task_getCurrent());
+  // task.  This satisfies the precondition on cancel_unlocked().
+  _swift_taskGroup_cancel_unlocked(group, swift_task_getCurrent());
 }
 
 // =============================================================================

stdlib/public/Concurrency/TaskPrivate.h

@@ -89,16 +89,16 @@ AsyncTask *_swift_task_clearCurrent();
 /// Set the active task reference for the current thread.
 AsyncTask *_swift_task_setCurrent(AsyncTask *newTask);
 
-/// Cancel all the child tasks that belong to `group`.
+/// Cancel the task group and all the child tasks that belong to `group`.
 ///
 /// The caller must guarantee that this is called while holding the owning
 /// task's status record lock.
-void _swift_taskGroup_cancelAllChildren(TaskGroup *group);
+void _swift_taskGroup_cancel(TaskGroup *group);
 
-/// Cancel all the child tasks that belong to `group`.
+/// Cancel the task group and all the child tasks that belong to `group`.
 ///
 /// The caller must guarantee that this is called from the owning task.
-void _swift_taskGroup_cancelAllChildren_unlocked(TaskGroup *group,
+void _swift_taskGroup_cancel_unlocked(TaskGroup *group,
                                                  AsyncTask *owningTask);
 
 /// Remove the given task from the given task group.

stdlib/public/Concurrency/TaskStatus.cpp

@@ -776,11 +776,13 @@ void swift::_swift_taskGroup_detachChild(TaskGroup *group,
   });
 }
 
-/// Cancel all the child tasks that belong to `group`.
+/// Cancel the task group and all the child tasks that belong to `group`.
 ///
 /// The caller must guarantee that this is called while holding the owning
 /// task's status record lock.
-void swift::_swift_taskGroup_cancelAllChildren(TaskGroup *group) {
+void swift::_swift_taskGroup_cancel(TaskGroup *group) {
+  (void) group->statusCancel();
+
   // Because only the owning task of the task group can modify the
   // child list of a task group status record, and it can only do so
   // while holding the owning task's status record lock, we do not need
@@ -789,10 +791,10 @@ void swift::_swift_taskGroup_cancelAllChildren(TaskGroup *group) {
     swift_task_cancel(childTask);
 }
 
-/// Cancel all the child tasks that belong to `group`.
+/// Cancel the task group and all the child tasks that belong to `group`.
 ///
 /// The caller must guarantee that this is called from the owning task.
-void swift::_swift_taskGroup_cancelAllChildren_unlocked(TaskGroup *group,
+void swift::_swift_taskGroup_cancel_unlocked(TaskGroup *group,
                                                         AsyncTask *owningTask) {
   // Early out. If there are no children, there's nothing to do. We can safely
   // check this without locking, since this can only be concurrently mutated
@@ -801,7 +803,7 @@ void swift::_swift_taskGroup_cancelAllChildren_unlocked(TaskGroup *group,
     return;
 
   withStatusRecordLock(owningTask, [&group](ActiveTaskStatus status) {
-    _swift_taskGroup_cancelAllChildren(group);
+    _swift_taskGroup_cancel(group);
   });
 }
 
@@ -825,7 +827,7 @@ static void performCancellationAction(TaskStatusRecord *record) {
   // under the synchronous control of the task that owns the group.
   case TaskStatusRecordKind::TaskGroup: {
     auto groupRecord = cast<TaskGroupTaskStatusRecord>(record);
-    _swift_taskGroup_cancelAllChildren(groupRecord->getGroup());
+    _swift_taskGroup_cancel(groupRecord->getGroup());
     return;
   }
 

test/Concurrency/Runtime/async_taskgroup_addUnlessCancelled.swift

@@ -0,0 +1,59 @@
+// RUN: %target-run-simple-swift( -target %target-swift-5.1-abi-triple -parse-as-library) | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+
+@available(SwiftStdlib 6.0, *)
+func test_withTaskGroup_addUnlessCancelled() async throws {
+  let task = Task {
+    await withTaskGroup(of: Void.self) { group in
+      print("Inner: Sleep...")
+      try? await Task.sleep(for: .seconds(60)) // we'll never actually wait 10 seconds, as this will be woken up by cancel
+      print("Inner: Task.isCancelled: \(Task.isCancelled)")
+
+      let added = group.addTaskUnlessCancelled {
+        print("Added Task! Child Task.isCancelled: \(Task.isCancelled)")
+      }
+      print("Inner: Task added = \(added)")  // CHECK: Inner: Task added = false
+    }
+  }
+
+  print("Outer: Cancel!")
+  task.cancel()
+  print("Outer: Cancelled")
+
+  await task.value
+}
+
+@available(SwiftStdlib 6.0, *)
+func test_withDiscardingTaskGroup_addUnlessCancelled() async throws {
+  let task = Task {
+    await withDiscardingTaskGroup { group in
+      print("Inner: Sleep...")
+      try? await Task.sleep(for: .seconds(60)) // we'll never actually wait 10 seconds, as this will be woken up by cancel
+      print("Inner: Task.isCancelled: \(Task.isCancelled)")
+
+      let added = group.addTaskUnlessCancelled {
+        print("Added Task! Child Task.isCancelled: \(Task.isCancelled)")
+      }
+      print("Inner: Task added = \(added)")  // CHECK: Inner: Task added = false
+    }
+  }
+
+  print("Outer: Cancel!")
+  task.cancel()
+  print("Outer: Cancelled")
+
+  await task.value
+}
+
+@available(SwiftStdlib 6.0, *)
+@main struct Main {
+  static func main() async {
+    try! await test_withTaskGroup_addUnlessCancelled()
+    try! await test_withDiscardingTaskGroup_addUnlessCancelled()
+  }
+}