Make sure that cancelling a task group does not invoke cancellation

handler of parent task that created the group

Change comment in TaskGroup.swift to enforce that only parent task can
call cancelAll on the group

Add tests to verify mutating of task group in child tasks will fail

Radar-Id: rdar://problem/86346865

Author: rokhinip

include/swift/ABI/TaskGroup.h

@@ -18,6 +18,7 @@
 #define SWIFT_ABI_TASK_GROUP_H
 
 #include "swift/ABI/Task.h"
+#include "swift/ABI/TaskStatus.h"
 #include "swift/ABI/HeapObject.h"
 #include "swift/Runtime/Concurrency.h"
 #include "swift/Runtime/Config.h"
@@ -46,6 +47,9 @@ class alignas(Alignment_TaskGroup) TaskGroup {
   // Add a child task to the group. Always called with the status record lock of
   // the parent task held
   void addChildTask(AsyncTask *task);
+
+  // Provide accessor for task group's status record
+  TaskGroupTaskStatusRecord *getTaskRecord();
 };
 
 } // end namespace swift

stdlib/public/Concurrency/TaskGroup.cpp

@@ -449,6 +449,10 @@ static TaskGroup *asAbstract(TaskGroupImpl *group) {
   return reinterpret_cast<TaskGroup*>(group);
 }
 
+TaskGroupTaskStatusRecord * TaskGroup::getTaskRecord() {
+    return asImpl(this)->getTaskRecord();
+}
+
 // =============================================================================
 // ==== initialize -------------------------------------------------------------
 

stdlib/public/Concurrency/TaskGroup.swift

@@ -373,8 +373,8 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
   /// If you add a task to a group after canceling the group,
   /// that task is canceled immediately after being added to the group.
   ///
-  /// There are no restrictions on where you can call this method.
-  /// Code inside a child task or even another task can cancel a group.
+  /// This method can only be called by the parent task that created the task
+  /// group.
   ///
   /// - SeeAlso: `Task.isCancelled`
   /// - SeeAlso: `TaskGroup.isCancelled`

stdlib/public/Concurrency/TaskStatus.cpp

@@ -485,41 +485,6 @@ static void performCancellationAction(TaskStatusRecord *record) {
   // FIXME: allow dynamic extension/correction?
 }
 
-/// Perform any cancellation actions required by the given record.
-static void performGroupCancellationAction(TaskStatusRecord *record) {
-  switch (record->getKind()) {
-  // We only need to cancel specific GroupChildTasks, not arbitrary child tasks.
-  // A task may be parent to many tasks which are not part of a group after all.
-  case TaskStatusRecordKind::ChildTask:
-    return;
-
-  case TaskStatusRecordKind::TaskGroup: {
-    auto groupChildRecord = cast<TaskGroupTaskStatusRecord>(record);
-    // Since a task can only be running a single task group at the same time,
-    // we can always assume that the group record which we found is the one
-    // we're intended to cancel child tasks for.
-    //
-    // A group enforces that tasks can not "escape" it, and as such once the group
-    // returns, all its task have been completed.
-    for (AsyncTask *child: groupChildRecord->children()) {
-      swift_task_cancel(child);
-    }
-    return;
-  }
-
-  // All other kinds of records we handle the same way as in a normal cancellation
-  case TaskStatusRecordKind::Deadline:
-  case TaskStatusRecordKind::CancellationNotification:
-  case TaskStatusRecordKind::EscalationNotification:
-  case TaskStatusRecordKind::Private_RecordLock:
-    performCancellationAction(record);
-    return;
-}
-
-  // Other cases can fall through here and be ignored.
-  // FIXME: allow dynamic extension/correction?
-}
-
 SWIFT_CC(swift)
 static void swift_task_cancelImpl(AsyncTask *task) {
   SWIFT_TASK_DEBUG_LOG("cancel task = %p", task);
@@ -543,17 +508,15 @@ SWIFT_CC(swift)
 static void swift_task_cancel_group_child_tasksImpl(TaskGroup *group) {
   // Acquire the status record lock.
   //
-  // We purposefully DO NOT make this a cancellation by itself.
-  // We are cancelling the task group, and all tasks it contains.
-  // We are NOT cancelling the entire parent task though.
+  // Guaranteed to be called from the context of the parent task that created
+  // the task group once we have #40616
   auto task = swift_task_getCurrent();
   withStatusRecordLock(task, LockContext::OnTask,
                        [&](ActiveTaskStatus &status) {
-    // Carry out the cancellation operations associated with all
-    // the active records.
-    for (auto cur: status.records()) {
-      performGroupCancellationAction(cur);
-    }
+    // We purposefully DO NOT make this a cancellation by itself.
+    // We are cancelling the task group, and all tasks it contains.
+    // We are NOT cancelling the entire parent task though.
+    performCancellationAction(group->getTaskRecord());
   });
 }
 

test/Concurrency/Runtime/async_taskgroup_cancel_from_inside_child.swift

@@ -0,0 +1,40 @@
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking %import-libdispatch -parse-as-library) | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: libdispatch
+
+// rdar://76038845
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+
+func test_taskGroup_cancelAll() async {
+
+    await withTaskCancellationHandler {
+        await withTaskGroup(of: Int.self, returning: Void.self) { group in
+            group.spawn {
+                await Task.sleep(3_000_000_000)
+                let c = Task.isCancelled
+                print("group task isCancelled: \(c)")
+                return 0
+           }
+
+           group.cancelAll() // Cancels the group but not the task
+           _ = await group.next()
+       }
+    } onCancel : {
+       print("parent task cancel handler called")
+    }
+
+    // CHECK-NOT: parent task cancel handler called
+    // CHECK: group task isCancelled: true
+    // CHECK: done
+    print("done")
+}
+
+@available(SwiftStdlib 5.1, *)
+@main struct Main {
+  static func main() async {
+    await test_taskGroup_cancelAll()
+  }
+}

test/Concurrency/Runtime/taskgroup_cancelAll_cancellationHandler.swift

@@ -0,0 +1,27 @@
+// RUN: %target-typecheck-verify-swift  -disable-availability-checking
+// REQUIRES: concurrency
+
+@available(SwiftStdlib 5.1, *)
+func test_taskGroup_cancelAll() async {
+    await withTaskGroup(of: Int.self, returning: Void.self) { group in
+        group.spawn {
+            await Task.sleep(3_000_000_000)
+            let c = Task.isCancelled
+            print("group task isCancelled: \(c)")
+            return 0
+       }
+
+       group.spawn {
+         group.cancelAll() //expected-warning{{capture of 'group' with non-sendable type 'TaskGroup<Int>' in a `@Sendable` closure}}
+                           //expected-error@-1{{reference to captured parameter 'group' in concurrently-executing code}}
+         return 0
+       }
+       group.spawn { [group] in
+         group.cancelAll() //expected-warning{{capture of 'group' with non-sendable type 'TaskGroup<Int>' in a `@Sendable` closure}}
+         return 0
+       }
+       _ = await group.next()
+    }
+
+    print("done")
+}