[swift-evolution] Contextualizing async coroutines

[Joe Groff]

The coroutine proposal as it stands essentially exposes raw delimited continuations. While this is a flexible and expressive feature in the abstract, for the concrete purpose of representing asynchronous coroutines, it provides weak user-level guarantees about where their code might be running after being resumed from suspension, and puts a lot of pressure on APIs to be well-behaved in this respect. And if we're building toward actors, where async actor methods should be guaranteed to run "in the actor", I think we'll *need* something more than the bare-bones delimited continuation approach to get there. I think the proposal's desire to keep coroutines independent of a specific runtime model is a good idea, but I also think there are a couple possible modifications we could add to the design to make it easier to reason about what context things run in for any runtime model that benefits from async/await:

# Coroutine context

Associating a context value with a coroutine would let us thread useful information through the execution of the coroutine. This is particularly useful for GCD, so you could attach a queue, QoS, and other attributes to the coroutine, since these aren't reliably available from the global environment. It could be a performance improvement even for things like per-pthread queues, since coroutine context should be cheaper to access than pthread_self. 

For example, a coroutine-aware `dispatch_async` could spawn a coroutine with the queue object and other interesting attributes as its context:

extension DispatchQueue {
 func `async`(_ body: () async -> ()) {
   dispatch_async(self, {
     beginAsync(context: self) { await body() }
   })
 }
}

and well-behaved dispatch-aware async APIs could use the context to decide how they should schedule completion:

func asyncOverlay() async -> T {
 // If the coroutine is associated with a queue, schedule completion on that queue
 if let currentQueue = getCoroutineContext() as? DispatchQueue {
   if #available(iOS 23, macOS 10.24, *) {
     // Maybe Apple frameworks add APIs that let you control completion dispatch up front…
     suspendAsync { continuation in
       originalAPI(on: currentQueue, completion: continuation)
     }
   } else {
     // …but we still need to queue-hop explicitly for backward deployment
     suspendAsync { continuation in
       originalAPI(completion: { dispatch_async(currentQueue, continuation) })
     }
   }
 } else {
   // If the coroutine isn't associated with a queue, leave it up to the API to schedule continuation
   suspendAsync { continuation in
     originalAPI(completion: continuation)
   }
 }
}

Similarly, if you've built your own framework on another platform with per-thread event loops and want to maintain thread affinity for coroutines through your APIs, you could similarly provide APIs that beginAsync with a thread ID as context and use that context to figure out where to schedule the continuation when you do something that suspends.

# `onResume` hooks

Relying on coroutine context alone still leaves responsibility wholly on suspending APIs to pay attention to the coroutine context and schedule the continuation correctly. You'd still have the expression problem when coroutine-spawning APIs from one framework interact with suspending APIs from another framework that doesn't understand the spawning framework's desired scheduling policy. We could provide some defense against this by letting the coroutine control its own resumption with an "onResume" hook, which would run when a suspended continuation is invoked instead of immediately resuming the coroutine. That would let the coroutine-aware dispatch_async example from above do something like this, to ensure the continuation always ends up back on the correct queue:

extension DispatchQueue {
 func `async`(_ body: () async -> ()) {
   dispatch_async(self, {
     beginAsync(
       context: self,
       body: { await body() },
       onResume: { continuation in
         // Defensively hop to the right queue
         dispatch_async(self, continuation)
       }
     )
   })
 }
}

This would let spawning APIs provide a stronger guarantee that the spawned coroutine is always executing as if scheduled by a specific queue/actor/event loop/HWND/etc., even if later suspended by an async API working in a different paradigm. This would also let you more strongly associate a coroutine with a future object representing its completion:

class CoroutineFuture<T> {
 enum State {
   case busy // currently running
   case suspended(() -> ()) // suspended
   case success(T) // completed with success
   case failure(Error) // completed with error
 }

 var state: State = .busy

 init(_ body: () async -> T) {

   beginAsync(
     body: {
       do {
         self.state = .success(await body())
       } catch {
         self.state = .failure(error)
       }
     },
     onResume: { continuation in
       assert(self.state == .busy, "already running?!")
       self.state = .suspended(continuation)
     }
   }
 }

 // Return the result of the future, or try to make progress computing it
 func poll() throws -> T? {
   switch state {
   case .busy:
     return nil
   case .suspended(let cont):
     cont()
     switch state {
     case .success(let value):
       return value
     case .failure(let error):
       throw error
     case .busy, .suspended:
       return nil
     }
   case .success(let value):
     return value
   case .error(let error):
     throw error
 }
}


A downside of this design is that it incurs some cost from defensive rescheduling on the continuation side, and also prevents writing APIs that intentionally change context across an `await`, like a theoretical "goToMainThread()" function (though you could do that by spawning a semantically-independent coroutine associated with the main thread, which might be a better design anyway).

-Joe