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

Seems that during refactorings of child cancellations we somehow missed
also cancelling the group itself. It seems we did not have good test
coverage of the addTaskUnlessCancelled somehow and thus this slipped
through.

This adds a regression test for addTaskUnlessCancelled and fixes how we
handle the cancellation effect in TaskStatus.

resolves swiftlang#80789
resolves rdar://149177600

Author: ktoso

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 ------------------------------------------------------------------
 

stdlib/public/Concurrency/TaskStatus.cpp

@@ -781,6 +781,7 @@ void swift::_swift_taskGroup_detachChild(TaskGroup *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) {
+  assert(group->isCancelled() && "Expected task group to be cancelled when cancelling all child tasks.");
   // 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
@@ -825,7 +826,12 @@ 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());
+    auto group = groupRecord->getGroup();
+    auto wasAlreadyCancelled = group->statusCancel();
+    if (wasAlreadyCancelled) {
+      return;
+    }
+    _swift_taskGroup_cancelAllChildren(group);
     return;
   }
 

test/Concurrency/Runtime/async_taskgroup_addUnlessCancelled.swift

@@ -0,0 +1,64 @@
+// RUN: %target-run-simple-swift( -target %target-swift-5.1-abi-triple -parse-as-library) | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+
+// rdar://76038845
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+
+import Dispatch
+
+@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(nanoseconds: 2_000_000_000)
+      print("Inner: Task.isCancelled: \(Task.isCancelled)")
+
+      let added = group.addTaskUnlessCancelled {
+        print("Added Task! Child Task.isCancelled: \(Task.isCancelled)")
+      }
+      print("Inner: Task added = \(added)")  // CHECK: Task added = false
+    }
+  }
+
+  try? await Task.sleep(nanoseconds: 1_000_000)
+  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(nanoseconds: 2_000_000_000)
+      print("Inner: Task.isCancelled: \(Task.isCancelled)")
+
+      let added = group.addTaskUnlessCancelled {
+        print("Added Task! Child Task.isCancelled: \(Task.isCancelled)")
+      }
+      print("Inner: Task added = \(added)")  // CHECK: Task added = false
+    }
+  }
+
+  try? await Task.sleep(nanoseconds: 1_000_000)
+  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()
+  }
+}