+concurrency task groups initial work in progress

Author: DougGregor

stdlib/public/Concurrency/Task.swift

@@ -148,7 +148,7 @@ extension Task {
   /// This is a port of the C++ FlagSet.
   struct JobFlags {
     /// Kinds of schedulable jobs.
-    enum Kind : Int {
+    enum Kind: Int {
       case task = 0
     };
 
@@ -215,54 +215,6 @@ extension Task {
 // ==== Detached Tasks ---------------------------------------------------------
 
 extension Task {
-  /// Run given `operation` as part of a new top-level task.
-  ///
-  /// Creating detached tasks should, generally, be avoided in favor of using
-  /// `async` functions, `async let` declarations and `await` expressions - as
-  /// those benefit from structured, bounded concurrency which is easier to reason
-  /// about, as well as automatically inheriting the parent tasks priority,
-  /// task-local storage, deadlines, as well as being cancelled automatically
-  /// when their parent task is cancelled. Detached tasks do not get any of those
-  /// benefits, and thus should only be used when an operation is impossible to
-  /// be modelled with child tasks.
-  ///
-  /// ### Cancellation
-  /// A detached task always runs to completion unless it is explicitly cancelled.
-  /// Specifically, dropping a detached tasks `Task.Handle` does _not_ automatically
-  /// cancel given task.
-  ///
-  /// Canceling a task must be performed explicitly via `handle.cancel()`.
-  ///
-  /// - Note: it is generally preferable to use child tasks rather than detached
-  ///   tasks. Child tasks automatically carry priorities, task-local state,
-  ///   deadlines and have other benefits resulting from the structured
-  ///   concurrency concepts that they model. Consider using detached tasks only
-  ///   when strictly necessary and impossible to model operations otherwise.
-  ///
-  /// - Parameters:
-  ///   - priority: priority of the task TODO: reword and define more explicitly once we have priorities well-defined
-  ///   - operation: the operation to execute
-  /// - Returns: handle to the task, allowing to `await handle.get()` on the
-  ///     tasks result or `cancel` it.
-  public static func runDetached<T>(
-    priority: Priority = .default,
-    operation: @escaping () async -> T
-  ) -> Handle<T> {
-    // Set up the job flags for a new task.
-    var flags = JobFlags()
-    flags.kind = .task
-    flags.priority = priority
-    flags.isFuture = true
-
-    // Create the asynchronous task future.
-    let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, nil, operation)
-
-    // Enqueue the resulting job.
-    _enqueueJobGlobal(Builtin.convertTaskToJob(task))
-
-    return Handle<T>(task: task)
-  }
-
   /// Run given throwing `operation` as part of a new top-level task.
   ///
   /// Creating detached tasks should, generally, be avoided in favor of using
@@ -288,7 +240,7 @@ extension Task {
   ///   when strictly necessary and impossible to model operations otherwise.
   ///
   /// - Parameters:
-  ///   - priority: priority of the task TODO: reword and define more explicitly once we have priorities well-defined
+  ///   - priority: priority of the task
   ///   - operation: the operation to execute
   /// - Returns: handle to the task, allowing to `await handle.get()` on the
   ///     tasks result or `cancel` it. If the operation fails the handle will
@@ -405,6 +357,19 @@ public func _runChildTask<T>(operation: @escaping () async throws -> T) async
   return task
 }
 
+//@_silgen_name("swift_task_get_priority") // TODO: not quite this way?
+//public func getTaskPriority(_ task: __owned Builtin.NativeObject) -> JobPriority
+
+struct RawTaskFutureWaitResult {
+  let hadErrorResult: Bool
+  let storage: UnsafeRawPointer
+}
+
+@_silgen_name("swift_task_future_wait")
+func taskFutureWait(
+  on task: Builtin.NativeObject
+) async -> RawTaskFutureWaitResult
+
 #if _runtime(_ObjC)
 
 /// Intrinsic used by SILGen to launch a task for bridging a Swift async method

stdlib/public/Concurrency/TaskGroup.swift

@@ -11,6 +11,7 @@
 ////===----------------------------------------------------------------------===//
 
 import Swift
+import Dispatch
 @_implementationOnly import _SwiftConcurrencyShims
 
 // ==== Task Group -------------------------------------------------------------
@@ -32,6 +33,12 @@ extension Task {
   ///       // some accumulation logic (e.g. sum += result)
   ///     }
   ///
+  /// ### Thrown errors
+  /// When tasks are added to the group using the `group.add` function, they may
+  /// immediately begin executing. Even if their results are not collected explicitly
+  /// and such task throws, and was not yet cancelled, it may result in the `withGroup`
+  /// throwing.
+  ///
   /// ### Cancellation
   /// If an error is thrown out of the task group, all of its remaining tasks
   /// will be cancelled and the `withGroup` call will rethrow that error.
@@ -57,61 +64,149 @@ extension Task {
   public static func withGroup<TaskResult, BodyResult>(
     resultType: TaskResult.Type,
     returning returnType: BodyResult.Type = BodyResult.self,
-    body: (inout Task.Group<TaskResult>) async throws -> BodyResult
-  ) async rethrows -> BodyResult {
-    fatalError("\(#function) not implemented yet.")
+    cancelOutstandingTasksOnReturn: Bool = false,
+    body: @escaping ((inout Task.Group<TaskResult>) async throws -> BodyResult)
+  ) async throws -> BodyResult {
+    let drainPendingTasksOnSuccessfulReturn = !cancelOutstandingTasksOnReturn
+    let parent = Builtin.getCurrentAsyncTask()
+
+    // Set up the job flags for a new task.
+    var groupFlags = JobFlags()
+    groupFlags.kind = .task // TODO: .taskGroup?
+    groupFlags.priority = .default // TODO: parent's priority // await Task.currentPriority()
+    groupFlags.isFuture = true
+
+    // 1. Prepare the Group task
+    // FIXME: do we have to rather prepare it inside the task we spawn; and yield it back along with the result instead?
+    var group = Task.Group<TaskResult>(parentTask: parent)
+
+    let (groupTask, context) =
+      Builtin.createAsyncTaskFuture(groupFlags.bits, nil) { () async throws -> BodyResult in
+        let got = await try body(&group)
+        return got
+      }
+    let groupHandle = Handle<BodyResult>(task: groupTask)
+
+    // 2.0) Run the task!
+    DispatchQueue.global(priority: .default).async {
+      groupHandle.run() // TODO: this synchronously runs
+    }
+
+    // 2.1) ensure that if we fail and exit by throwing we will cancel all tasks,
+    // if we succeed, there is nothing to cancel anymore so this is noop
+    defer { group.cancelAll() }
+
+    // 2.2) Await the group completing it's run ("until the withGroup returns")
+    let result = await try groupHandle.get() // if we throw, so be it -- group tasks will be cancelled
+
+// TODO: do drain before exiting
+//    if drainPendingTasksOnSuccessfulReturn {
+//      // drain all outstanding tasks
+//      while await try group.next() != nil {
+//        continue // awaiting all remaining tasks
+//      }
+//    }
+
+    return result
   }
 
   /// A task group serves as storage for dynamically started tasks.
   ///
-  /// Its intended use is with the
+  /// Its intended use is with the `Task.withGroup` function.
   /* @unmoveable */
   public struct Group<TaskResult> {
+    private let parentTask: Builtin.NativeObject
+
+    // TODO: we want groups to be unordered in completion, the counterpart to streams (Series),
+    // as such it feels like we need to keep them like this, because a next() can complete any of them
+    // and then we need to remove it from the pending ones
+    private var pendingTasks: [Int: Handle<TaskResult>] // TODO: make a dict for out of order completions
+    private var nextTaskID: Int = 0
+
+    /// If present, the handle on which the `next()` call is awaiting,
+    /// it should be resumed by *any* of the in-flight tasks completing.
+    private var nextHandle: Task.Handle<TaskResult>? = nil
+
     /// No public initializers
-    private init() {}
+    init(parentTask: Builtin.NativeObject) {
+      self.parentTask = parentTask
+      self.pendingTasks = [:]
+    }
 
     // Swift will statically prevent this type from being copied or moved.
     // For now, that implies that it cannot be used with generics.
 
     /// Add a child task to the group.
     ///
     /// ### Error handling
-    /// Operations are allowed to throw.
-    ///
-    /// in which case the `await try next()`
+    /// Operations are allowed to `throw`, in which case the `await try next()`
     /// invocation corresponding to the failed task will re-throw the given task.
     ///
+    /// The `add` function will never (re)-throw exceptions from the `operation`,
+    /// the corresponding `next()` call will throw the error when necessary.
+    ///
     /// - Parameters:
     ///   - overridingPriority: override priority of the operation task
     ///   - operation: operation to execute and add to the group
+    @discardableResult
     public mutating func add(
       overridingPriority: Priority? = nil,
-      operation: () async throws -> TaskResult
-    ) async {
-      fatalError("\(#function) not implemented yet.")
-    }
+      operation: @escaping () async throws -> TaskResult
+    ) async -> Task.Handle<TaskResult> {
+      var flags = JobFlags()
+      flags.kind = .task // TODO: childTask?
+      flags.priority = .default // TODO: priority getting from parent not implemented yet
+//      if let overridingPriority = overridingPriority { // TODO: cannot use ?? with async defaultValue
+//        flags.priority = overridingPriority
+//      } else {
+//        flags.priority = await Task.currentPriority() // TODO: self.parent.priority ?
+//      }
+      flags.isFuture = true
 
-    /// Add a child task and return a `Task.Handle` that can be used to manage it.
-    ///
-    /// The task's result is accessible either via the returned `handle` or the
-    /// `group.next()` function (as any other `add`-ed task).
-    ///
-    /// - Parameters:
-    ///   - overridingPriority: override priority of the operation task
-    ///   - operation: operation to execute and add to the group
-    public mutating func addWithHandle(
-      overridingPriority: Priority? = nil,
-      operation: () async throws -> TaskResult
-    ) async -> Handle<TaskResult> {
-      fatalError("\(#function) not implemented yet.")
+      let (childTask, context) =
+        // TODO: passing the parentTask (instead of nil) here makes the program hang here
+        Builtin.createAsyncTaskFuture(flags.bits, nil, operation)
+
+      let handle = Handle<TaskResult>(task: childTask)
+
+      // runTask(childTask)
+        DispatchQueue.global(priority: .default).async {
+          handle.run()
+        }
+
+//      _ = DispatchQueue.global(priority: .default).async {
+//          print("run dispatch INSIDE: \(childTask)")
+//          await try operation()
+//        }
+
+      // FIXME: need to store? self.pendingTasks[ObjectIdentifier(childTask)] = childTask
+
+      defer { nextTaskID += 1 }
+      self.pendingTasks[nextTaskID] = handle
+
+      return handle
     }
 
     /// Wait for a child task to complete and return the result it returned,
     /// or else return.
     ///
-    ///
-    public mutating func next() async throws -> TaskResult? {
-      fatalError("\(#function) not implemented yet.")
+    /// Order of completions is *not* guaranteed to be same as submission order,
+    /// rather the order of `next()` calls completing is by completion order of
+    /// the tasks. This differentiates task groups from streams (
+    public mutating func next(file: String = #file, line: UInt = #line) async throws -> TaskResult? {
+      // FIXME: this implementation is wrong and naive; we instead need to maintain a dict of handles,
+      //        and return them as they complete in that order; so likely a queue of "which one completed"
+      //        this will allow building "collect first N results" APIs easily;
+      //        APIs which need order can implement on top of this, or we provide a different API for it
+      let handle = self.pendingTasks.removeValue(forKey: 0) ??
+        self.pendingTasks.removeValue(forKey: 1)
+
+      if let handle = handle {
+        let got = await try handle.get()
+        return got
+      } else {
+        return nil
+      }
     }
 
     /// Query whether the group has any remaining tasks.
@@ -122,7 +217,7 @@ extension Task {
     ///
     /// - Returns: `true` if the group has no pending tasks, `false` otherwise.
     public var isEmpty: Bool {
-      fatalError("\(#function) not implemented yet.")
+      return self.pendingTasks.isEmpty
     }
 
     /// Cancel all the remaining tasks in the group.
@@ -133,8 +228,11 @@ extension Task {
     /// cancellation, are silently discarded.
     ///
     /// - SeeAlso: `Task.addCancellationHandler`
-    public mutating func cancelAll() {
-      fatalError("\(#function) not implemented yet.")
+    public mutating func cancelAll(file: String = #file, line: UInt = #line) {
+      for (id, handle) in self.pendingTasks {
+        handle.cancel()
+      }
+      self.pendingTasks = [:]
     }
   }
 }

test/Concurrency/Runtime/async_task_group_basic.swift

@@ -0,0 +1,74 @@
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency) | %FileCheck %s --dump-input always
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: OS=macosx
+
+import Dispatch
+
+// ==== ------------------------------------------------------------------------
+// MARK: "Infrastructure" for the tests
+
+extension DispatchQueue {
+  func async<R>(operation: @escaping () async -> R) -> Task.Handle<R> {
+    let handle = Task.runDetached(operation: operation)
+
+    // Run the task
+    _ = { self.async { handle.run() } }() // force invoking the non-async version
+
+    return handle
+  }
+}
+
+@available(*, deprecated, message: "This is a temporary hack")
+func launch<R>(operation: @escaping () async -> R) -> Task.Handle<R> {
+  let handle = Task.runDetached(operation: operation)
+
+  // Run the task
+  _ = DispatchQueue.global(priority: .default).async { handle.run() }
+
+  return handle
+}
+
+// ==== ------------------------------------------------------------------------
+// MARK: Tests
+
+func test_taskGroup_01_sum() {
+  let numbers = [1, 2]
+  let expected = numbers.reduce(0, +)
+
+  let taskHandle = launch { () async -> Int in
+    return await try! Task.withGroup(resultType: Int.self) { (group) async -> Int in
+      for n in numbers {
+        await group.add { () async -> Int in n }
+      }
+
+      var sum = 0
+      while let r = await try! group.next() {
+        print("next: \(r)")
+        sum += r
+      }
+
+      print("task group returning: \(sum)")
+      return sum
+    }
+  }
+
+  // CHECK: main task
+  // CHECK: next: 1
+  // CHECK: next: 2
+  // CHECK: task group returning: 3
+
+  launch { () async in
+    let sum = await try! taskHandle.get()
+    // CHECK: result: 3
+    print("result: \(sum)")
+    assert(expected == sum)
+    exit(0)
+  }
+
+  print("main task")
+}
+
+test_taskGroup_01_sum()
+
+dispatchMain()