[Concurrency][TaskGroup] allow cancelAll be invoked from child tasks

Author: ktoso

include/swift/ABI/TaskGroup.h

@@ -51,17 +51,12 @@ namespace swift {
         Error = 0b11,
     };
 
-    /// Describes the status of the waiting task that is suspended on `next()`.
-    enum class WaitStatus : uintptr_t {
-        Waiting = 0,
-    };
-
     enum class GroupPollStatus : uintptr_t {
         /// The channel is known to be empty and we can immediately return nil.
         Empty = 0,
 
         /// The task has been enqueued to the channels wait queue.
-        Waiting = 1,
+        MustWait = 1,
 
         /// The task has completed with result (of type \c resultType).
         Success = 2,
@@ -331,12 +326,17 @@ namespace swift {
       return s;
     }
 
+    /// Add a single pending task to the status counter.
+    /// This is used to implement next() properly, as we need to know if there
+    /// are pending tasks worth suspending/waiting for or not.
+    ///
+    /// Note that the group does *not* store child tasks at all, as they are
+    /// stored in the `TaskGroupTaskStatusRecord` inside the current task, that
+    /// is currently executing the group. Here we only need the counts of
+    /// pending/ready tasks.
+    ///
     /// Returns *assumed* new status, including the just performed +1.
-    GroupStatus statusAddPendingTaskRelaxed(
-//        AsyncTask* pendingTask
-        ) {
-//      assert(pendingTask->isFuture());
-
+    GroupStatus statusAddPendingTaskRelaxed() {
       auto old = status.fetch_add(GroupStatus::onePendingTask, std::memory_order_relaxed);
       auto s = GroupStatus {old + GroupStatus::onePendingTask };
 
@@ -346,30 +346,13 @@ namespace swift {
         s = GroupStatus {o - GroupStatus::onePendingTask };
       }
 
-      fprintf(stderr, "[%s:%d] (%s): status %s\n", __FILE__, __LINE__, __FUNCTION__, s.to_string().c_str());
-
-//      // FIXME: we won't need the +1 in the status, just the queue?
-//      pendingQueue.enqueue(PendingQueueItem::get(pendingTask))
       return s;
     }
 
     GroupStatus statusLoadRelaxed() {
       return GroupStatus{status.load(std::memory_order_relaxed)};
     }
 
-//    /// Returns *assumed* new status, including the just performed +1.
-//    GroupStatus statusAddWaitingTaskAcquire() {
-//      auto old = status.fetch_add(GroupStatus::oneWaitingTask, std::memory_order_acquire);
-//      return GroupStatus { old + GroupStatus::oneWaitingTask };
-//    }
-
-//    /// Remove waiting task, without taking any pending task.
-//    GroupStatus statusRemoveWaitingTask() {
-//      return GroupStatus {
-//          status.fetch_sub(GroupStatus::oneWaitingTask, std::memory_order_relaxed)
-//      };
-//    }
-
     /// Compare-and-set old status to a status derived from the old one,
     /// by simultaneously decrementing one Pending and one Waiting tasks.
     ///
@@ -397,7 +380,7 @@ namespace swift {
     /// If unable to complete the waiting task immediately (with an readily
     /// available completed task), either returns an `GroupPollStatus::Empty`
     /// result if it is known that no pending tasks in the group,
-    /// or a `GroupPollStatus::Waiting` result if there are tasks in flight
+    /// or a `GroupPollStatus::MustWait` result if there are tasks in flight
     /// and the waitingTask eventually be woken up by a completion.
     TaskGroup::PollResult poll(AsyncTask *waitingTask);
 

stdlib/public/Concurrency/Task.cpp

@@ -111,9 +111,7 @@ void AsyncTask::completeFuture(AsyncContext *context, ExecutorRef executor) {
   assert(queueHead.getStatus() == Status::Executing);
 
   // If this is task group child, notify the parent group about the completion.
-  fprintf(stderr, "[%s:%d] (%s) complete task %d\n", __FILE__, __LINE__, __FUNCTION__, this);
   if (hasGroupChildFragment()) {
-    fprintf(stderr, "[%s:%d] (%s) complete group child task %d\n", __FILE__, __LINE__, __FUNCTION__, this);
     // then we must offer into the parent group that we completed,
     // so it may `next()` poll completed child tasks in completion order.
     auto group = groupChildFragment()->getGroup();
@@ -141,7 +139,6 @@ static void destroyTask(SWIFT_CONTEXT HeapObject *obj) {
 
   // For a future, destroy the result.
   if (task->isFuture()) {
-    fprintf(stderr, "[%s:%d] (%s): task: %d\n", __FILE__, __LINE__, __FUNCTION__, task);
     task->futureFragment()->destroy();
   }
 
@@ -313,13 +310,7 @@ AsyncTaskAndContext swift::swift_task_create_group_future_f(
 
   // Perform additional linking between parent and child task.
   if (parent) {
-    // FIXME: we must attach children (`async let` created child tasks to the parent)
-//    if (!flags.task_isGroupChildTask()) {
-//      // just a normal child task
-//      swift_task_attachChild(parent, task); // TODO: this has to be done outside of here (!!!!!!!!!!!!!!!!!!!)
-//    } // else, group children are recorded outside
-
-    // if the parent was already cancelled, we carry this flag forward to the child.
+    // If the parent was already cancelled, we carry this flag forward to the child.
     //
     // 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,

stdlib/public/Concurrency/TaskGroup.cpp

@@ -91,13 +91,6 @@ void TaskGroup::destroy(AsyncTask *task) {
   swift_task_removeStatusRecord(task, Record);
   swift_task_dealloc(task, Record);
 
-  // TODO: need to release all waiters as well
-//  auto waitHead = waitQueue.load(std::memory_order_acquire);
-//  switch (waitHead.getStatus()) {
-//    case TaskGroup::WaitStatus::Waiting:
-//      assert(false && "destroying a task group that still has waiting tasks");
-//  }
-
   mutex.lock(); // TODO: remove fragment lock, and use status for synchronization
   // Release all ready tasks which are kept retained, the group destroyed,
   // so no other task will ever await on them anymore;
@@ -223,7 +216,7 @@ void swift::swift_task_group_wait_next(
   TaskGroup::PollResult polled = group->poll(waitingTask);
   fprintf(stderr, "[%s:%d] (%s): group polled: %d\n", __FILE__, __LINE__, __FUNCTION__, polled.status);
 
-  if (polled.status == TaskGroup::GroupPollStatus::Waiting) {
+  if (polled.status == TaskGroup::GroupPollStatus::MustWait) {
     fprintf(stderr, "[%s:%d] (%s): group polled: WAITING\n", __FILE__, __LINE__, __FUNCTION__);
     // The waiting task has been queued on the channel,
     // there were pending tasks so it will be woken up eventually.
@@ -280,7 +273,7 @@ TaskGroup::PollResult TaskGroup::poll(AsyncTask *waitingTask) {
       ReadyQueueItem item;
       bool taskDequeued = readyQueue.dequeue(item);
       if (!taskDequeued) {
-        result.status = TaskGroup::GroupPollStatus::Waiting;
+        result.status = TaskGroup::GroupPollStatus::MustWait;
         result.storage = nullptr;
         result.retainedTask = nullptr;
         mutex.unlock(); // TODO: remove group lock, and use status for synchronization
@@ -336,7 +329,7 @@ TaskGroup::PollResult TaskGroup::poll(AsyncTask *waitingTask) {
       fprintf(stderr, "[%s:%d] (%s): added to wait queue\n", __FILE__, __LINE__, __FUNCTION__);
       mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
       // no ready tasks, so we must wait.
-      result.status = TaskGroup::GroupPollStatus::Waiting;
+      result.status = TaskGroup::GroupPollStatus::MustWait;
       return result;
     } // else, try again
   }

stdlib/public/Concurrency/TaskGroup.swift

@@ -13,12 +13,6 @@
 import Swift
 @_implementationOnly import _SwiftConcurrencyShims
 
-//import Darwin // FIXME: remove this
-
-func pprint(_ m: String, file: String = #file, line: UInt = #line) {
-//  fputs("[\(file):\(line)] \(m)\n", stderr)
-}
-
 // ==== Task Group -------------------------------------------------------------
 
 extension Task {
@@ -69,7 +63,6 @@ extension Task {
   public static func withGroup<TaskResult, BodyResult>(
     resultType: TaskResult.Type,
     returning returnType: BodyResult.Type = BodyResult.self,
-//    startingChildTasksOn executor: ExecutorRef? = nil, // TODO: actually respect it
     body: (inout Task.Group<TaskResult>) async throws -> BodyResult
   ) async throws -> BodyResult {
     let task = Builtin.getCurrentAsyncTask()
@@ -184,6 +177,12 @@ extension Task {
     /// Awaiting on an empty group results in the immediate return of a `nil`
     /// value, without the group task having to suspend.
     ///
+    /// It is also possible to use `for await` to collect results of a task groups:
+    ///
+    ///     for await try value in group {
+    ///         collected += value
+    ///     }
+    ///
     /// ### Thread-safety
     /// Please note that the `group` object MUST NOT escape into another task.
     /// The `group.next()` MUST be awaited from the task that had originally
@@ -263,10 +262,13 @@ extension Task {
     /// Any results, including errors thrown by tasks affected by this
     /// cancellation, are silently discarded.
     ///
+    /// This function may be called even from within child (or any other) tasks,
+    /// and will reliably cause the group to become cancelled.
+    ///
     /// - SeeAlso: `Task.addCancellationHandler`
     /// - SeeAlso: `Task.checkCancelled`
     /// - SeeAlso: `Task.isCancelled`
-    public mutating func cancelAll() {
+    public func cancelAll() {
       _taskGroupCancelAll(task: _task, group: _group)
     }
 

stdlib/public/Concurrency/TaskPrivate.h

@@ -143,8 +143,8 @@ static void runTaskWithPollResult(
       // return a `nil` here (as result of the `group.next()`)
       waitingTaskContext->result.storage = nullptr;
       break;
-    case TaskGroup::GroupPollStatus::Waiting:
-      assert(false && "Must not attempt to run with a Waiting result.");
+    case TaskGroup::GroupPollStatus::MustWait:
+      assert(false && "Must not attempt to run with a MustWait result.");
   }
 
   // FIXME: removing children on next() resumption is not implemented yet; we keep accumulating them until the group exits

test/Concurrency/Runtime/async_task_cancellation_early.swift

@@ -4,6 +4,8 @@
 // REQUIRES: concurrency
 // REQUIRES: libdispatch
 
+import Dispatch
+
 #if canImport(Darwin)
 import Darwin
 #elseif canImport(Glibc)
@@ -14,7 +16,7 @@ func test_runDetached_cancel_child_early() async {
   print(#function) // CHECK: test_runDetached_cancel_child_early
   let h: Task.Handle<Bool, Error> = Task.runDetached {
     async let childCancelled: Bool = { () -> Bool in
-      usleep(1000 * 2)
+      sleep(2)
       return await Task.__unsafeCurrentAsync().isCancelled
     }()
 

test/Concurrency/Runtime/async_task_cancellation_while_running.swift

@@ -4,6 +4,8 @@
 // REQUIRES: concurrency
 // REQUIRES: libdispatch
 
+import Dispatch
+
 #if canImport(Darwin)
 import Darwin
 #elseif canImport(Glibc)
@@ -13,7 +15,7 @@ import Glibc
 func test_runDetached_cancel_while_child_running() async {
   let h: Task.Handle<Bool, Error> = Task.runDetached {
     async let childCancelled: Bool = { () -> Bool in
-      usleep(1000 * 3)
+      sleep(3)
       return await Task.__unsafeCurrentAsync().isCancelled
     }()
 
@@ -25,7 +27,7 @@ func test_runDetached_cancel_while_child_running() async {
   }
 
   // sleep here, i.e. give the task a moment to start running
-  usleep(1000 * 2)
+  sleep(2)
 
   h.cancel()
   print("handle cancel")

test/Concurrency/Runtime/async_task_priority_current.swift

@@ -1,6 +1,10 @@
 // RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck --dump-input=always %s
+
 // REQUIRES: executable_test
 // REQUIRES: concurrency
+// REQUIRES: libdispatch
+
+import Dispatch
 
 #if canImport(Darwin)
 import Darwin
@@ -30,7 +34,7 @@ func test_multiple_lo_indirectly_escalated() async {
   @concurrent
   func loopUntil(priority: Task.Priority) async {
     while (await Task.__unsafeCurrentAsync().task.priority != priority) {
-      usleep(1000 * 1)
+      sleep(1)
     }
   }
 

test/Concurrency/Runtime/async_taskgroup_cancelAll_only_specific_group.swift

@@ -1,6 +1,10 @@
 // RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input always
+
 // REQUIRES: executable_test
 // REQUIRES: concurrency
+// REQUIRES: libdispatch
+
+import Dispatch
 
 #if canImport(Darwin)
 import Darwin
@@ -24,7 +28,7 @@ func test_taskGroup_cancelAll_onlySpecificGroup() async {
 
     for i in 1...5 {
       await group.add {
-        usleep(1000 * 1)
+        sleep(1)
         let c = await Task.__unsafeCurrentAsync().isCancelled
         pprint("add: \(i) (cancelled: \(c))")
         return i
@@ -49,7 +53,7 @@ func test_taskGroup_cancelAll_onlySpecificGroup() async {
   let g2: Int = try! await Task.withGroup(resultType: Int.self) { group in
     for i in 1...3 {
       await group.add {
-        usleep(1000 * 1)
+        sleep(1)
         let c = await Task.__unsafeCurrentAsync().isCancelled
         pprint("g1 task \(i) (cancelled: \(c))")
         return i

test/Concurrency/Runtime/async_taskgroup_cancel_from_inside_child.swift

@@ -0,0 +1,44 @@
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input always
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: libdispatch
+
+import Dispatch
+
+#if canImport(Darwin)
+import Darwin
+#elseif canImport(Glibc)
+import Glibc
+#endif
+
+func test_taskGroup_cancel_from_inside_child() async {
+  let result: Int = try! await Task.withGroup(resultType: Int.self) { group in
+    let firstAdded = await group.add { [group] in // must explicitly capture, as the task executes concurrently
+      group.cancelAll() // allowed
+      print("first")
+      return 1
+    }
+    print("firstAdded: \(firstAdded)") // CHECK: firstAdded: true
+
+    let one = try! await group.next()
+
+    let secondAdded = await group.add {
+      print("second")
+      return 2
+    }
+    print("secondAdded: \(secondAdded)") // CHECK: secondAdded: false
+
+    return 1
+  }
+
+  print("result: \(result)") // CHECK: result: 1
+}
+
+
+
+@main struct Main {
+  static func main() async {
+    await test_taskGroup_cancel_from_inside_child()
+  }
+}