Merge pull request swiftlang#36307 from ktoso/wip-yield

[Concurrency] Simple Task.yield implementation

Author: ktoso

stdlib/public/Concurrency/Task.swift

@@ -503,6 +503,37 @@ extension Task {
   }
 }
 
+// ==== Voluntary Suspension -----------------------------------------------------
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+extension Task {
+
+  /// Explicitly suspend the current task, potentially giving up execution actor
+  /// of current actor/task, allowing other tasks to execute.
+  ///
+  /// This is not a perfect cure for starvation;
+  /// if the task is the highest-priority task in the system, it might go
+  /// immediately back to executing.
+  public static func yield() async {
+    // Prepare the job flags
+    var flags = JobFlags()
+    flags.kind = .task
+    flags.priority = .default
+    flags.isFuture = true
+
+    // Create the asynchronous task future, it will do nothing, but simply serves
+    // as a way for us to yield our execution until the executor gets to it and
+    // resumes us.
+    // TODO: consider if it would be useful for this task to be a child task
+    let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, {})
+
+    // Enqueue the resulting job.
+    _enqueueJobGlobal(Builtin.convertTaskToJob(task))
+
+    let _ = await Handle<Void, Never>(task).get()
+  }
+}
+
 // ==== UnsafeCurrentTask ------------------------------------------------------
 
 /// Calls the given closure with the with the "current" task in which this

stdlib/public/Concurrency/TaskGroup.swift

@@ -708,7 +708,7 @@ extension TaskGroup: AsyncSequence {
     /// - SeeAlso: `TaskGroup.next()` for a detailed discussion its semantics.
     public mutating func next() async -> Element? {
       guard !finished else { return nil }
-      guard let element = try await group.next() else {
+      guard let element = await group.next() else {
         finished = true
         return nil
       }

test/Concurrency/Runtime/async_task_yield.swift

@@ -0,0 +1,49 @@
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency %import-libdispatch -parse-as-library) | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+protocol Go: Actor {
+  func go(times: Int) async -> Int
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+extension Go {
+  func go(times: Int) async -> Int {
+    for i in 0...times {
+      print("\(Self.self) @ \(i)")
+      await Task.yield()
+    }
+    return times
+  }
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+actor One: Go {}
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+actor Two: Go {}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+func yielding() async {
+  let one = One()
+  let two = Two()
+  await withTaskGroup(of: Int.self) { group in
+    await group.spawn {
+      await one.go(times: 100)
+    }
+    await group.spawn {
+      await two.go(times: 100)
+    }
+  }
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+@main struct Main {
+  static func main() async {
+    await yielding()
+    // TODO: No idea for a good test for this... Open to ideas?
+    // CHECK:      One
+    // CHECK-NEXT: Two
+  }
+}