refactor: modify task operation to asynchronous wait to use stream

Author: soumyamahunt

Sources/AsyncObjects/AsyncSemaphore.swift

@@ -1,5 +1,4 @@
 import Foundation
-
 import OrderedCollections
 
 /// An object that controls access to a resource across multiple task contexts through use of a traditional counting semaphore.

Sources/AsyncObjects/TaskOperation.swift

@@ -28,8 +28,8 @@ import Dispatch
 /// // or wait synchronously for completion
 /// operation.waitUntilFinished()
 /// ```
-public final class TaskOperation<R: Sendable>: Operation, AsyncObject,
-    ContinuableCollection, Loggable, @unchecked Sendable
+public final class TaskOperation<R: Sendable>: Operation, AsyncObject, Loggable,
+    @unchecked Sendable
 {
     /// The asynchronous action to perform as part of the operation..
     private let underlyingAction: @Sendable () async throws -> R
@@ -40,6 +40,12 @@ public final class TaskOperation<R: Sendable>: Operation, AsyncObject,
     /// synchronize data access and modifications.
     @usableFromInline
     internal let locker: Locker
+    /// The stream that propagates the operation
+    /// finish event.
+    internal var event: AsyncStream<Void>!
+    /// The continuation of `stream` that controls the asynchronous wait
+    /// for operation completion.
+    internal var waiter: AsyncStream<Void>.Continuation!
 
     /// A type representing a set of behaviors for the executed
     /// task type and task completion behavior.
@@ -116,9 +122,8 @@ public final class TaskOperation<R: Sendable>: Operation, AsyncObject,
             willChangeValue(forKey: "isFinished")
             locker.perform {
                 _isFinished = newValue
-                guard newValue, !continuations.isEmpty else { return }
-                continuations.forEach { $1.resume() }
-                continuations = [:]
+                guard newValue else { return }
+                waiter.finish()
             }
             didChangeValue(forKey: "isFinished")
         }
@@ -166,11 +171,14 @@ public final class TaskOperation<R: Sendable>: Operation, AsyncObject,
         self.flags = flags
         self.underlyingAction = operation
         super.init()
+        self.event = AsyncStream(
+            bufferingPolicy: .bufferingOldest(1)
+        ) { self.waiter = $0 }
     }
 
     deinit {
         execTask?.cancel()
-        locker.perform { self.continuations.forEach { $1.cancel() } }
+        waiter.finish()
     }
 
     /// Begins the execution of the operation.
@@ -222,99 +230,6 @@ public final class TaskOperation<R: Sendable>: Operation, AsyncObject,
     }
 
     // MARK: AsyncObject
-    /// The suspended tasks continuation type.
-    @usableFromInline
-    internal typealias Continuation = TrackedContinuation<
-        GlobalContinuation<Void, Error>
-    >
-    /// The continuations stored with an associated key for all the suspended task that are waiting for operation completion.
-    @usableFromInline
-    internal private(set) var continuations: [UUID: Continuation] = [:]
-
-    /// Add continuation with the provided key in `continuations` map.
-    ///
-    /// - Parameters:
-    ///   - continuation: The `continuation` to add.
-    ///   - key: The key in the map.
-    ///   - file: The file add request originates from (there's usually no need to pass it
-    ///           explicitly as it defaults to `#fileID`).
-    ///   - function: The function add request originates from (there's usually no need to
-    ///               pass it explicitly as it defaults to `#function`).
-    ///   - line: The line add request originates from (there's usually no need to pass it
-    ///           explicitly as it defaults to `#line`).
-    ///   - preinit: The pre-initialization handler to run
-    ///              in the beginning of this method.
-    ///
-    /// - Important: The pre-initialization handler must run
-    ///              before any logic in this method.
-    @inlinable
-    internal func addContinuation(
-        _ continuation: Continuation,
-        withKey key: UUID,
-        file: String, function: String, line: UInt,
-        preinit: @Sendable () -> Void
-    ) {
-        locker.perform {
-            preinit()
-            log("Adding", id: key, file: file, function: function, line: line)
-            guard !continuation.resumed else {
-                log(
-                    "Already resumed, not tracking", id: key,
-                    file: file, function: function, line: line
-                )
-                return
-            }
-
-            guard !isFinished else {
-                continuation.resume()
-                log(
-                    "Resumed", id: key,
-                    file: file, function: function, line: line
-                )
-                return
-            }
-
-            continuations[key] = continuation
-            log("Tracking", id: key, file: file, function: function, line: line)
-        }
-    }
-
-    /// Remove continuation associated with provided key
-    /// from `continuations` map.
-    ///
-    /// - Parameters:
-    ///   - continuation: The continuation to remove and cancel.
-    ///   - key: The key in the map.
-    ///   - file: The file remove request originates from (there's usually no need to pass it
-    ///           explicitly as it defaults to `#fileID`).
-    ///   - function: The function remove request originates from (there's usually no need to
-    ///               pass it explicitly as it defaults to `#function`).
-    ///   - line: The line remove request originates from (there's usually no need to pass it
-    ///           explicitly as it defaults to `#line`).
-    @inlinable
-    internal func removeContinuation(
-        _ continuation: Continuation,
-        withKey key: UUID,
-        file: String, function: String, line: UInt
-    ) {
-        locker.perform {
-            log("Removing", id: key, file: file, function: function, line: line)
-            continuations.removeValue(forKey: key)
-            guard !continuation.resumed else {
-                log(
-                    "Already resumed, not cancelling", id: key,
-                    file: file, function: function, line: line
-                )
-                return
-            }
-
-            continuation.cancel()
-            log(
-                "Cancelled", id: key,
-                file: file, function: function, line: line
-            )
-        }
-    }
 
     /// Starts operation asynchronously
     /// as part of a new top-level task on behalf of the current actor.
@@ -356,18 +271,28 @@ public final class TaskOperation<R: Sendable>: Operation, AsyncObject,
         function: String = #function,
         line: UInt = #line
     ) async throws {
-        guard !isFinished else {
-            log("Finished", file: file, function: function, line: line)
-            return
-        }
-
         let key = UUID()
         log("Waiting", id: key, file: file, function: function, line: line)
-        try await withPromisedContinuation(
-            withKey: key,
-            file: file, function: function, line: line
-        )
-        log("Finished", id: key, file: file, function: function, line: line)
+        for await _ in event { break }
+        do {
+            try Task.checkCancellation()
+        } catch {
+            log(
+                "Cancelled", id: key,
+                file: file, function: function, line: line
+            )
+            throw error
+        }
+        do {
+            let _ = try await execTask?.value
+            log("Finished", id: key, file: file, function: function, line: line)
+        } catch {
+            log(
+                "Finished with error: \(error)", id: key,
+                file: file, function: function, line: line
+            )
+            throw error
+        }
     }
 }
 

Sources/AsyncObjects/TaskQueue.swift

@@ -1,5 +1,4 @@
 import Foundation
-
 import OrderedCollections
 
 /// An object that acts as a concurrent queue executing submitted tasks concurrently.