Merge pull request #40562 from apple/rokhinip/86346865-taskgroup-cancellAll-cancellationHandler

Cancellation of TaskGroup should not trigger cancel handler of parent task

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

@@ -457,6 +457,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")
+}