Make ActorQueue's tasks isolated to an actor's context

Author: dfed

Sources/AsyncQueue/ActorQueue.swift

@@ -31,14 +31,14 @@
 ///
 ///     nonisolated
 ///     public func log(_ message: String) {
-///         queue.async {
-///             await self.append(message)
+///         queue.async(on: self) { myself in
+///             myself.append(message)
 ///         }
 ///     }
 ///
 ///     nonisolated
 ///     public func retrieveLogs() async -> [String] {
-///         await queue.await { await self.logs }
+///         await queue.await(on: self) { myself in myself.logs }
 ///     }
 ///
 ///     private func append(_ message: String) {
@@ -82,32 +82,38 @@ public final class ActorQueue {
 
     /// 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)
+    /// - Parameters:
+    ///   - isolatedActor: The actor within which the task is isolated.
+    ///   - task: The task to enqueue.
+    public func async<ActorType: Actor>(on isolatedActor: ActorType, _ task: @escaping @Sendable (isolated ActorType) async -> Void) {
+        taskStreamContinuation.yield { await task(isolatedActor) }
     }
 
     /// 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.
+    /// - Parameters:
+    ///   - isolatedActor: The actor within which the task is isolated.
+    ///   - task: The task to enqueue.
     /// - Returns: The value returned from the enqueued task.
-    public func await<T>(_ task: @escaping @Sendable () async -> T) async -> T {
+    public func await<ActorType: Actor, T>(on isolatedActor: isolated ActorType, _ task: @escaping @Sendable (isolated ActorType) async -> T) async -> T {
         await withUnsafeContinuation { continuation in
             taskStreamContinuation.yield {
-                continuation.resume(returning: await task())
+                continuation.resume(returning: await task(isolatedActor))
             }
         }
     }
 
     /// 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.
+    /// - Parameters:
+    ///   - isolatedActor: The actor within which the task is isolated.
+    ///   - task: The task to enqueue.
     /// - Returns: The value returned from the enqueued task.
-    public func await<T>(_ task: @escaping @Sendable () async throws -> T) async throws -> T {
+    public func await<ActorType: Actor, T>(on isolatedActor: isolated ActorType, _ task: @escaping @Sendable (isolated ActorType) async throws -> T) async throws -> T {
         try await withUnsafeThrowingContinuation { continuation in
             taskStreamContinuation.yield {
                 do {
-                    continuation.resume(returning: try await task())
+                    continuation.resume(returning: try await task(isolatedActor))
                 } catch {
                     continuation.resume(throwing: error)
                 }

Tests/AsyncQueueTests/ActorQueueTests.swift

@@ -39,47 +39,46 @@ final class ActorQueueTests: XCTestCase {
     func test_async_sendsEventsInOrder() async {
         let counter = Counter()
         for iteration in 1...1_000 {
-            systemUnderTest.async {
-                await counter.incrementAndExpectCount(equals: iteration)
+            systemUnderTest.async(on: counter) { counter in
+                counter.incrementAndExpectCount(equals: iteration)
             }
         }
-        await systemUnderTest.await { /* Drain the queue */ }
+        await systemUnderTest.await(on: counter) { _ in /* Drain the queue */ }
     }
 
     func test_async_startsExecutionOfNextTaskAfterSuspension() async {
         let semaphore = Semaphore()
-        systemUnderTest.async {
+        systemUnderTest.async(on: semaphore) { semaphore in
             await semaphore.wait()
         }
-        systemUnderTest.async {
+        systemUnderTest.async(on: semaphore) { semaphore in
             // Signal the semaphore from the actor queue.
             // If the actor queue were FIFO, this test would hang since this code would never execute:
             // we'd still be waiting for the prior `wait()` tasks to finish.
-            await semaphore.signal()
+            semaphore.signal()
         }
-        await systemUnderTest.await { /* Drain the queue */ }
+        await systemUnderTest.await(on: semaphore) { _ in /* Drain the queue */ }
     }
 
     func test_await_allowsReentrancy() async {
         let counter = Counter()
-        await systemUnderTest.await { [systemUnderTest] in
-            await systemUnderTest.await {
-                await counter.incrementAndExpectCount(equals: 1)
+        await systemUnderTest.await(on: counter) { [systemUnderTest] counter in
+            await systemUnderTest.await(on: counter) { counter in
+                counter.incrementAndExpectCount(equals: 1)
             }
-            await counter.incrementAndExpectCount(equals: 2)
+            counter.incrementAndExpectCount(equals: 2)
         }
-        await systemUnderTest.await { /* Drain the queue */ }
     }
 
     func test_async_executesEnqueuedTasksAfterReceiverIsDeallocated() async {
         var systemUnderTest: ActorQueue? = ActorQueue()
         let counter = Counter()
         let expectation = self.expectation(description: #function)
         let semaphore = Semaphore()
-        systemUnderTest?.async {
+        systemUnderTest?.async(on: counter) { counter in
             // Make the task wait.
             await semaphore.wait()
-            await counter.incrementAndExpectCount(equals: 1)
+            counter.incrementAndExpectCount(equals: 1)
             expectation.fulfill()
         }
         weak var queue = systemUnderTest
@@ -110,14 +109,14 @@ final class ActorQueueTests: XCTestCase {
         let asyncSemaphore = Semaphore()
         let syncSemaphore = Semaphore()
         let expectation = self.expectation(description: #function)
-        systemUnderTest.async { [reference = referenceHolder.reference] in
+        systemUnderTest.async(on: syncSemaphore) { [reference = referenceHolder.reference] syncSemaphore in
             // Now that we've started the task and captured the reference, release the synchronous code.
-            await syncSemaphore.signal()
+            syncSemaphore.signal()
             // Wait for the synchronous setup to complete and the reference to be nil'd out.
             await asyncSemaphore.wait()
             // Retain the unsafe counter until the task is completed.
             _ = reference
-            self.systemUnderTest.async {
+            self.systemUnderTest.async(on: syncSemaphore) { _ in
                 // Signal that this task has cleaned up.
                 // This closure will not execute until the prior closure completes.
                 expectation.fulfill()
@@ -138,26 +137,26 @@ final class ActorQueueTests: XCTestCase {
     func test_await_sendsEventsInOrder() async {
         let counter = Counter()
         for iteration in 1...1_000 {
-            systemUnderTest.async {
-                await counter.incrementAndExpectCount(equals: iteration)
+            systemUnderTest.async(on: counter) { counter in
+                counter.incrementAndExpectCount(equals: iteration)
             }
 
             guard iteration % 25 == 0 else {
                 // Keep sending async events to the queue.
                 continue
             }
 
-            await systemUnderTest.await {
-                let count = await counter.count
+            await systemUnderTest.await(on: counter) { counter in
+                let count = counter.count
                 XCTAssertEqual(count, iteration)
             }
         }
-        await systemUnderTest.await { /* Drain the queue */ }
+        await systemUnderTest.await(on: counter) { counter in /* Drain the queue */ }
     }
 
     func test_await_canReturn() async {
         let expectedValue = UUID()
-        let returnedValue = await systemUnderTest.await { expectedValue }
+        let returnedValue = await systemUnderTest.await(on: Counter()) { _ in expectedValue }
         XCTAssertEqual(expectedValue, returnedValue)
     }
 
@@ -167,7 +166,7 @@ final class ActorQueueTests: XCTestCase {
         }
         let expectedError = TestError()
         do {
-            try await systemUnderTest.await { throw expectedError }
+            try await systemUnderTest.await(on: Counter()) { _ in throw expectedError }
         } catch {
             XCTAssertEqual(error as? TestError, expectedError)
         }