async_task_groups.swift

// RUN: %target-swift-frontend -target %target-swift-5.1-abi-triple %s -emit-sil -o /dev/null -verify
// RUN: %target-swift-frontend -target %target-swift-5.1-abi-triple %s -emit-sil -o /dev/null -verify -strict-concurrency=targeted
// RUN: %target-swift-frontend -target %target-swift-5.1-abi-triple %s -emit-sil -o /dev/null -verify -strict-concurrency=complete -verify-additional-prefix tns-

// REQUIRES: concurrency
// REQUIRES: asserts
// REQUIRES: libdispatch

@available(SwiftStdlib 5.1,*)
func asyncFunc()async->Int{42}
@available(SwiftStdlib 5.1,*)
func asyncThrowsFunc()asyncthrows->Int{42}
@available(SwiftStdlib 5.1,*)
func asyncThrowsOnCancel()asyncthrows->Int{
 // terrible suspend-spin-loop -- do not do this
 // only for purposes of demonstration
whileTask.isCancelled {
try?awaitTask.sleep(nanoseconds:1_000_000_000)
}

throwCancellationError()
}

@available(SwiftStdlib 5.1,*)
func test_taskGroup_add()asyncthrows->Int{
tryawaitwithThrowingTaskGroup(of:Int.self){ group in
 group.addTask{
awaitasyncFunc()
}

 group.addTask{
awaitasyncFunc()
}

varsum=0
whilelet v =tryawait group.next(){
 sum += v
}
 return sum
} // implicitly awaits
}

// ==== ------------------------------------------------------------------------
// MARK: Example group Usages

structBoom:Error{}
@available(SwiftStdlib 5.1,*)
func work()async->Int{42}
@available(SwiftStdlib 5.1,*)
func boom()asyncthrows->Int{throwBoom()}

@available(SwiftStdlib 5.1,*)
func first_allMustSucceed()asyncthrows{

letfirst:Int=tryawaitwithThrowingTaskGroup(of:Int.self){ group in
 group.addTask{awaitwork()}
 group.addTask{awaitwork()}
 group.addTask{tryawaitboom()}

iflet first =tryawait group.next(){
return first
} else {
fatalError("Should never happen, we either throw, or get a result from any of the tasks")
}
 // implicitly await: boom
}
 _ = first
 // Expected: re-thrown Boom
}

@available(SwiftStdlib 5.1,*)
func first_ignoreFailures()asyncthrows{
@Sendablefunc work()async->Int{42}
@Sendablefunc boom()asyncthrows->Int{throwBoom()}

letfirst:Int=tryawaitwithThrowingTaskGroup(of:Int.self){ group in
 group.addTask{awaitwork()}
 group.addTask{awaitwork()}
 group.addTask{
do{
returntryawaitboom()
}catch{
return0 // TODO: until try? await works properly
}
}

varresult:Int=0
whilelet v =tryawait group.next(){
 result = v

if result !=0{
break
}
}

 return result
}
 _ = first
 // Expected: re-thrown Boom
}

// ==== ------------------------------------------------------------------------
// MARK: Advanced Custom Task Group Usage

@available(SwiftStdlib 5.1,*)
func test_taskGroup_quorum_thenCancel()async{
 // imitates a typical "gather quorum" routine that is typical in distributed systems programming
enumVote{
case yay
case nay
}
structFollower:Sendable{
init(_ name:String){}
func vote()asyncthrows->Vote{
 // "randomly" vote yes or no
return.yay
}
}

 /// Performs a simple quorum vote among the followers.
 ///
 /// - Returns: `true` iff `N/2 + 1` followers return `.yay`, `false` otherwise.
func gatherQuorum(followers:[Follower])async->Bool{
try!await withThrowingTaskGroup(of:Vote.self){ group in
forfollowerin followers {
 group.addTask{tryawait follower.vote()}
}

defer{
 group.cancelAll()
}

varyays:Int=0
varnays:Int=0
letquorum=Int(followers.count /2)+1
whilelet vote =tryawait group.next(){
switch vote {
case.yay:
 yays +=1
if yays >= quorum {
 // cancel all remaining voters, we already reached quorum
returntrue
}
case.nay:
 nays +=1
if nays >= quorum {
returnfalse
}
}
}

 return false
}
}

 _ =awaitgatherQuorum(followers:[Follower("A"),Follower("B"),Follower("C")])
}

// FIXME: this is a workaround since (A, B) today isn't inferred to be Sendable
// and causes an error, but should be a warning (this year at least)
@available(SwiftStdlib 5.1,*)
struct SendableTuple2<A: Sendable, B: Sendable>: Sendable {
letfirst:A
letsecond:B

init(_ first: A, _ second: B){
self.first = first
self.second = second
}
}

@available(SwiftStdlib 5.1,*)
extension Collection where Self: Sendable, Element: Sendable,Self.Index: Sendable {

 /// Just another example of how one might use task groups.
func map<T:Sendable>(
 parallelism requestedParallelism:Int?=nil/*system default*/,
 // ordered: Bool = true, /
 _ transform:@Sendable(Element)asyncthrows->T // expected-note {{parameter 'transform' is implicitly non-escaping}}
)asyncthrows->[T]{ // TODO: can't use rethrows here, maybe that's just life though; rdar://71479187 (rethrows is a bit limiting with async functions that use task groups)
letdefaultParallelism=2
letparallelism= requestedParallelism ?? defaultParallelism

letn=self.count
if n ==0{
return[]
}

returntryawaitwithThrowingTaskGroup(of: SendableTuple2<Int, T>.self){ group in
varresult=ContiguousArray<T>()
 result.reserveCapacity(n)

vari=self.startIndex
varsubmitted=0

func submitNext()asyncthrows{
 // The reason that we emit an error here is b/c we capture the var box
 // to i and that is task isolated. This is the region isolation version
 // of the 'escaping closure captures non-escaping parameter' error.
 //
 // TODO: When we have isolation history, isolation history will be able
 // to tell us what is going on.
 group.addTask{[submitted,i]in // expected-error {{escaping closure captures non-escaping parameter 'transform'}}
let _ =tryawaittransform(self[i]) // expected-note {{captured here}}
letvalue:T?=nil
returnSendableTuple2(submitted, value!)
}
 submitted +=1
formIndex(after:&i)
}

 // submit first initial tasks
for_in0..<parallelism {
tryawaitsubmitNext()
}

whilelet tuple =tryawait group.next(){
letindex= tuple.first
lettaskResult= tuple.second
result[index]= taskResult

tryTask.checkCancellation()
tryawaitsubmitNext()
}

 assert(result.count == n)
return Array(result)
}
}
}