Make ActorQueue tasks execute on the target actor's execution context (#9)

* Make ActorQueue tasks execute on the target actor's execution context

* Documentation and naming pass

* Add comments to aid a consumer if they hit the unowned crash

* Review feedback

Co-authored-by: Michael Bachand <bachand.michael@gmail.com>

Author: dfed

AsyncQueue.podspec

@@ -1,6 +1,6 @@
 Pod::Spec.new do |s|
   s.name     = 'AsyncQueue'
-  s.version  = '0.1.0'
+  s.version  = '0.2.0'
   s.license  = 'MIT'
   s.summary  = 'A queue that enables ordered sending of events from synchronous to asynchronous code.'
   s.homepage = 'https://github.com/dfed/swift-async-queue'

README.md

@@ -16,21 +16,23 @@ Tasks sent from a synchronous context to an asynchronous context in Swift Concur
 @MainActor
 func testMainActorTaskOrdering() async {
     actor Counter {
-        func increment() -> Int {
+        func incrementAndAssertCountEquals(_ expectedCount: Int) {
             count += 1
-            return count
+            let incrementedCount = count
+            XCTAssertEqual(incrementedCount, expectedCount) // often fails
         }
-        var count = 0
+
+        private var count = 0
     }
 
     let counter = Counter()
     var tasks = [Task<Void, Never>]()
     for iteration in 1...100 {
         tasks.append(Task {
-            let incrementedCount = await counter.increment()
-            XCTAssertEqual(incrementedCount, iteration) // often fails
+            await counter.incrementAndAssertCountEquals(iteration)
         })
     }
+    // Wait for all enqueued tasks to finish.
     for task in tasks {
         _ = await task.value
     }
@@ -43,101 +45,87 @@ While [actors](https://docs.swift.org/swift-book/LanguageGuide/Concurrency.html#
 
 ### Executing asynchronous tasks in FIFO order
 
-Use a `FIFOQueue` to execute asynchronous tasks enqueued from a nonisolated context in FIFO order. Tasks sent to one of these queues are guaranteed to begin _and end_ executing in the order in which they are enqueued.
-
-```swift
-let queue = FIFOQueue()
-queue.async {
-    /*
-    `async` context that executes after all other enqueued work is completed.
-    Work enqueued after this task will wait for this task to complete.
-    */
-    try? await Task.sleep(nanoseconds: 1_000_000)
-}
-queue.async {
-    /*
-    This task begins execution once the above one-second sleep completes.
-    */
-}
-await queue.await {
-    /*
-    `async` context that can return a value or throw an error.
-    Executes after all other enqueued work is completed.
-    Work enqueued after this task will wait for this task to complete.
-    */
-}
-```
+Use a `FIFOQueue` to execute asynchronous tasks enqueued from a nonisolated context in FIFO order. Tasks sent to one of these queues are guaranteed to begin _and end_ executing in the order in which they are enqueued. A `FIFOQueue` executes tasks in a similar manner to a `DispatchQueue`: enqueued tasks executes atomically, and the program will deadlock if a task executing on a `FIFOQueue` awaits results from the queue on which it is executing.
 
-With a `FIFOQueue` you can easily execute asynchronous tasks from a nonisolated context in FIFO order:
+A `FIFOQueue` can easily execute asynchronous tasks from a nonisolated context in FIFO order:
 ```swift
 func testFIFOQueueOrdering() async {
     actor Counter {
-        func increment() -> Int {
+        nonisolated
+        func incrementAndAssertCountEquals(_ expectedCount: Int) {
+            queue.async {
+                await self.increment()
+                let incrementedCount = await self.count
+                XCTAssertEqual(incrementedCount, expectedCount) // always succeeds
+            }
+        }
+
+        nonisolated
+        func flushQueue() async {
+            await queue.await { }
+        }
+
+        func increment() {
             count += 1
-            return count
         }
+
         var count = 0
+
+        private let queue = FIFOQueue()
     }
 
     let counter = Counter()
-    let queue = FIFOQueue()
     for iteration in 1...100 {
-        queue.async {
-            let incrementedCount = await counter.increment()
-            XCTAssertEqual(incrementedCount, iteration) // always succeeds
-        }
+        counter.incrementAndAssertCountEquals(iteration)
     }
-    await queue.await { }
+    // Wait for all enqueued tasks to finish.
+    await counter.flushQueue()
 }
 ```
 
-### Sending ordered asynchronous tasks to Actors
+FIFO execution has a key downside: the queue must wait for all previously enqueued work – including suspended work – to complete before new work can begin. If you desire new work to start when a prior task suspends, utilize an `ActorQueue`.
 
-Use an `ActorQueue` to send ordered asynchronous tasks from a nonisolated context to an `actor` instance. Tasks sent to one of these queues are guaranteed to begin executing in the order in which they are enqueued. Ordering of execution is guaranteed up until the first [suspension point](https://docs.swift.org/swift-book/LanguageGuide/Concurrency.html#ID639) within the called `actor` code.
+### Sending ordered asynchronous tasks to Actors from a nonisolated context
 
-```swift
-let queue = ActorQueue()
-queue.async {
-    /*
-    `async` context that executes after all other enqueued work has begun executing.
-    Work enqueued after this task will wait for this task to complete or suspend.
-    */
-    try? await Task.sleep(nanoseconds: 1_000_000)
-}
-queue.async {
-    /*
-    This task begins execution once the above task suspends due to the one-second sleep.
-    */
-}
-await queue.await {
-    /*
-    `async` context that can return a value or throw an error.
-    Executes after all other enqueued work has begun executing.
-    Work enqueued after this task will wait for this task to complete or suspend.
-    */
-}
-```
+Use an `ActorQueue` to send ordered asynchronous tasks to an `actor`’s isolated context from nonisolated or synchronous contexts. Tasks sent to an actor queue are guaranteed to begin executing in the order in which they are enqueued. However, unlike a `FIFOQueue`, execution order is guaranteed only until the first [suspension point](https://docs.swift.org/swift-book/LanguageGuide/Concurrency.html#ID639) within the enqueued task. An `ActorQueue` executes tasks within the its adopted actor’s isolated context, resulting in `ActorQueue` task execution having the same properties as `actor` code execution: code between suspension points is executed atomically, and tasks sent to a single `ActorQueue` can await results from the queue without deadlocking.
+
+An instance of an `ActorQueue` is designed to be utilized by a single `actor` instance: tasks sent to an `ActorQueue` utilize the isolated context of the queue‘s adopted `actor` to serialize tasks. As such, there are a few requirements that must be met when dealing with an `ActorQueue`:
+1. The lifecycle of any `ActorQueue` should not exceed the lifecycle of its `actor`. It is strongly recommended that an `ActorQueue` be a `let` constant on the adopted `actor`. Enqueuing a task to an `ActorQueue` isntance after its adopted `actor` has been deallocated will result in a crash.
+2. An `actor` utilizing an `ActorQueue` should set the adopted execution context of the queue to `self` within the `actor`’s `init`. Failing to set an adopted execution context prior to enqueuing work on an `ActorQueue` will result in a crash.
 
-With an `ActorQueue` you can easily begin execution of asynchronous tasks from a nonisolated context in order:
+An `ActorQueue` can easily enqueue tasks that execute on an actor’s isolated context from a nonisolated context in order:
 ```swift
 func testActorQueueOrdering() async {
     actor Counter {
-        func increment() -> Int {
-            count += 1
-            return count
+        init() {
+            // Adopting the execution context in `init` satisfies requirement #2 above.
+            queue.adoptExecutionContext(of: self)
         }
-        var count = 0
+
+        nonisolated
+        func incrementAndAssertCountEquals(_ expectedCount: Int) {
+            queue.async { myself in
+                myself.count += 1
+                XCTAssertEqual(expectedCount, myself.count) // always succeeds
+            }
+        }
+
+        nonisolated
+        func flushQueue() async {
+            await queue.await { _ in }
+        }
+
+        private var count = 0
+        // Making the queue a private let constant satisfies requirement #1 above.
+        private let queue = ActorQueue<Counter>()
     }
 
     let counter = Counter()
-    let queue = ActorQueue()
     for iteration in 1...100 {
-        queue.async {
-            let incrementedCount = await counter.increment()
-            XCTAssertEqual(incrementedCount, iteration) // always succeeds
-        }
+        counter.incrementAndAssertCountEquals(iteration)
     }
-    await queue.await { }
+    // Wait for all enqueued tasks to finish.
+    await counter.flushQueue()
 }
 ```
 

Sources/AsyncQueue/ActorQueue.swift

@@ -20,56 +20,54 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
 
-/// A queue that executes asynchronous tasks enqueued from a nonisolated context.
+/// A queue that enables enqueing ordered asynchronous tasks from a nonisolated context onto an adopted actor's isolated context.
 /// Tasks are guaranteed to begin executing in the order in which they are enqueued. However, if a task suspends it will allow subsequently enqueued tasks to begin executing.
-/// Asynchronous tasks sent to this queue execute as they would in an `actor` type, allowing for re-entrancy and non-FIFO behavior when an individual task suspends.
+/// This queue exhibits the execution behavior of an actor: tasks sent to this queue can re-enter the queue, and tasks may execute in non-FIFO order when a task suspends.
 ///
-/// An `ActorQueue` is used to ensure tasks sent from a nonisolated context to a single `actor`'s isolated context begin execution in order.
-/// Here is an example of how an `ActorQueue` should be utilized within an `actor`:
+/// An `ActorQueue` ensures tasks sent from a nonisolated context to a single actor's isolated context begin execution in order.
+/// Here is an example of how an `ActorQueue` should be utilized within an actor:
 /// ```swift
 /// public actor LogStore {
 ///
+///     public init() {
+///         queue.adoptExecutionContext(of: self)
+///     }
+///
 ///     nonisolated
 ///     public func log(_ message: String) {
-///         queue.async {
-///             await self.append(message)
+///         queue.async { myself in
+///             myself.logs.append(message)
 ///         }
 ///     }
 ///
 ///     nonisolated
 ///     public func retrieveLogs() async -> [String] {
-///         await queue.await { await self.logs }
-///     }
-///
-///     private func append(_ message: String) {
-///         logs.append(message)
+///         await queue.await { myself in myself.logs }
 ///     }
 ///
-///     private let queue = ActorQueue()
+///     private let queue = ActorQueue<LogStore>()
 ///     private var logs = [String]()
 /// }
 /// ```
 ///
-/// - Warning: Execution order is not guaranteed unless the enqueued tasks interact with a single `actor` instance.
-public final class ActorQueue {
+/// - Precondition: The lifecycle of an `ActorQueue` must not exceed that of the adopted actor.
+public final class ActorQueue<ActorType: Actor> {
 
     // MARK: Initialization
 
     /// Instantiates an actor queue.
-    /// - Parameter priority: The baseline priority of the tasks added to the asynchronous queue.
-    public init(priority: TaskPriority? = nil) {
-        var capturedTaskStreamContinuation: AsyncStream<@Sendable () async -> Void>.Continuation? = nil
-        let taskStream = AsyncStream<@Sendable () async -> Void> { continuation in
+    public init() {
+        var capturedTaskStreamContinuation: AsyncStream<ActorTask>.Continuation? = nil
+        let taskStream = AsyncStream<ActorTask> { continuation in
             capturedTaskStreamContinuation = continuation
         }
         // Continuation will be captured during stream creation, so it is safe to force unwrap here.
         // If this force-unwrap fails, something is fundamentally broken in the Swift runtime.
         taskStreamContinuation = capturedTaskStreamContinuation!
 
-        Task.detached(priority: priority) {
-            let executor = ActorExecutor()
-            for await task in taskStream {
-                await executor.suspendUntilStarted(task)
+        Task.detached {
+            for await actorTask in taskStream {
+                await actorTask.executionContext.suspendUntilStarted(actorTask.task)
             }
         }
     }
@@ -80,65 +78,91 @@ public final class ActorQueue {
 
     // MARK: Public
 
+    /// Sets the actor context within which each `async` and `await`ed task will execute.
+    /// It is recommended that this method be called in the adopted actor’s `init` method.
+    /// **Must be called prior to enqueuing any work on the receiver.**
+    ///
+    /// - Parameter actor: The actor on which the queue's task will execute. This parameter is not retained by the receiver.
+    /// - Warning: Calling this method more than once will result in an assertion failure.
+    public func adoptExecutionContext(of actor: ActorType) {
+        assert(weakExecutionContext == nil) // Adopting multiple executionContexts on the same queue is API abuse.
+        weakExecutionContext = actor
+    }
+
     /// Schedules an asynchronous task for execution and immediately returns.
     /// The scheduled task will not execute until all prior tasks have completed or suspended.
-    /// - Parameter task: The task to enqueue.
-    public func async(_ task: @escaping @Sendable () async -> Void) {
-        taskStreamContinuation.yield(task)
+    /// - Parameter task: The task to enqueue. The task's parameter is a reference to the actor whose execution context has been adopted.
+    public func async(_ task: @escaping @Sendable (isolated ActorType) async -> Void) {
+        taskStreamContinuation.yield(ActorTask(executionContext: executionContext, task: task))
     }
 
     /// Schedules an asynchronous task and returns after the task is complete.
     /// The scheduled task will not execute until all prior tasks have completed or suspended.
-    /// - Parameter task: The task to enqueue.
+    /// - Parameter task: The task to enqueue. The task's parameter is a reference to the actor whose execution context has been adopted.
     /// - Returns: The value returned from the enqueued task.
-    public func await<T>(_ task: @escaping @Sendable () async -> T) async -> T {
-        await withUnsafeContinuation { continuation in
-            taskStreamContinuation.yield {
-                continuation.resume(returning: await task())
-            }
+    public func await<T>(_ task: @escaping @Sendable (isolated ActorType) async -> T) async -> T {
+        let executionContext = self.executionContext // Capture/retain the executionContext before suspending.
+        return await withUnsafeContinuation { continuation in
+            taskStreamContinuation.yield(ActorTask(executionContext: executionContext) { executionContext in
+                continuation.resume(returning: await task(executionContext))
+            })
         }
     }
 
     /// Schedules an asynchronous throwing task and returns after the task is complete.
     /// The scheduled task will not execute until all prior tasks have completed or suspended.
-    /// - Parameter task: The task to enqueue.
+    /// - Parameter task: The task to enqueue. The task's parameter is a reference to the actor whose execution context has been adopted.
     /// - Returns: The value returned from the enqueued task.
-    public func await<T>(_ task: @escaping @Sendable () async throws -> T) async throws -> T {
-        try await withUnsafeThrowingContinuation { continuation in
-            taskStreamContinuation.yield {
+    public func await<T>(_ task: @escaping @Sendable (isolated ActorType) async throws -> T) async throws -> T {
+        let executionContext = self.executionContext // Capture/retain the executionContext before suspending.
+        return try await withUnsafeThrowingContinuation { continuation in
+            taskStreamContinuation.yield(ActorTask(executionContext: executionContext) { executionContext in
                 do {
-                    continuation.resume(returning: try await task())
+                    continuation.resume(returning: try await task(executionContext))
                 } catch {
                     continuation.resume(throwing: error)
                 }
-            }
+            })
         }
     }
 
     // MARK: Private
 
-    private let taskStreamContinuation: AsyncStream<@Sendable () async -> Void>.Continuation
+    private let taskStreamContinuation: AsyncStream<ActorTask>.Continuation
 
-    // MARK: - ActorExecutor
+    /// The actor on whose isolated context our tasks run, force-unwrapped.
+    /// Utilize this accessor to retrieve the weak execution context in order to avoid repeating the below comment.
+    private var executionContext: ActorType {
+        // Crashing here means that this queue is being sent tasks either before an execution context has been set, or
+        // after the execution context has deallocated. An ActorQueue's execution context should be set in the adopted
+        // actor's `init` method, and the ActorQueue should not exceed the lifecycle of the adopted actor.
+        weakExecutionContext!
+    }
+    /// The actor on whose isolated context our tasks run.
+    /// We must use`weak` here to avoid creating a retain cycle between the adopted actor and this actor queue.
+    ///
+    /// We will assume this execution context always exists for the lifecycle of the queue because:
+    /// 1. The lifecycle of any `ActorQueue` must not exceed the lifecycle of its adopted `actor`.
+    /// 2. The adopted `actor` must set itself as the execution context for this queue within its `init` method.
+    private weak var weakExecutionContext: ActorType?
+
+    private struct ActorTask {
+        let executionContext: ActorType
+        let task: @Sendable (isolated ActorType) async -> Void
+    }
 
-    private actor ActorExecutor {
-        func suspendUntilStarted(_ task: @escaping @Sendable () async -> Void) async {
-            // Suspend the calling code until our enqueued task starts.
-            await withUnsafeContinuation { continuation in
-                // Utilize the serial (but not FIFO) Actor context to execute the task without requiring the calling method to wait for the task to complete.
-                Task {
-                    // Force this task to execute within the ActorExecutor's context by accessing an ivar on the instance.
-                    // This works around a bug when compiling with Xcode 14.1: https://github.com/apple/swift/issues/62503
-                    _ = void
+}
 
-                    // Signal that the task has started. As long as the `task` below interacts with another `actor` the order of execution is guaranteed.
-                    continuation.resume()
-                    await task()
-                }
+extension Actor {
+    func suspendUntilStarted(_ task: @escaping @Sendable (isolated Self) async -> Void) async {
+        // Suspend the calling code until our enqueued task starts.
+        await withUnsafeContinuation { continuation in
+            // Utilize the serial (but not FIFO) Actor context to execute the task without requiring the calling method to wait for the task to complete.
+            Task {
+                // Signal that the task has started. Since this `task` is executing on the current actor's execution context, the order of execution is guaranteed.
+                continuation.resume()
+                await task(self)
             }
         }
-
-        private let void: Void = ()
     }
-
 }