feat: add async semaphore

Author: soumyamahunt

.gitignore

@@ -37,14 +37,14 @@ playground.xcworkspace
 # Swift Package Manager
 #
 # Add this line if you want to avoid checking in source code from Swift Package Manager dependencies.
-# Packages/
-# Package.pins
-# Package.resolved
-# *.xcodeproj
-#
+Packages/
+Package.pins
+Package.resolved
+*.xcodeproj
+
 # Xcode automatically generates this directory with a .xcworkspacedata file and xcuserdata
 # hence it is not needed unless you have added a package configuration file to your project
-# .swiftpm
+.swiftpm
 
 .build/
 

Package.swift

@@ -0,0 +1,37 @@
+// swift-tools-version: 5.5
+
+import PackageDescription
+
+let package = Package(
+    name: "AsyncObject",
+    platforms: [
+        .macOS(.v10_15),
+        .iOS(.v13),
+        .tvOS(.v13),
+        .watchOS(.v6),
+    ],
+    products: [
+        .library(
+            name: "AsyncObject",
+            targets: ["AsyncObject"]
+        ),
+    ],
+    dependencies: [
+        .package(
+            url: "https://github.com/apple/swift-collections.git",
+            .upToNextMajor(from: "1.0.0")
+        ),
+    ],
+    targets: [
+        .target(
+            name: "AsyncObject",
+            dependencies: [
+                .product(name: "OrderedCollections", package: "swift-collections"),
+            ]
+        ),
+        .testTarget(
+            name: "AsyncObjectTests",
+            dependencies: ["AsyncObject"]
+        ),
+    ]
+)

Sources/AsyncObject/AsyncMutex.swift

@@ -0,0 +1,70 @@
+import Foundation
+
+public actor AsyncMutex {
+    private typealias Continuation = UnsafeContinuation<Void, Error>
+    private var continuations: [UUID: Continuation] = [:]
+    private var locked: Bool
+
+    private func addContinuation(
+        _ continuation: Continuation,
+        withKey key: UUID
+    ) {
+        continuations[key] = continuation
+    }
+
+    private func removeContinuation(withKey key: UUID) {
+        continuations.removeValue(forKey: key)
+    }
+
+    public init(lockedInitially locked: Bool = true) {
+        self.locked = locked
+    }
+
+    public func lock() {
+        locked = true
+    }
+
+    public func release() {
+        continuations.forEach { $0.value.resume() }
+        continuations = [:]
+        locked = false
+    }
+
+    public func wait() async {
+        guard locked else { return }
+        let key = UUID()
+        do {
+            try await withUnsafeThrowingContinuationCancellationHandler(
+                handler: { (continuation: Continuation) in
+                    Task { await removeContinuation(withKey: key) }
+                },
+                { addContinuation($0, withKey: key) }
+            )
+        } catch {
+            debugPrint(
+                "Wait on mutex for continuation task with key: \(key)"
+                + " cancelled with error \(error)"
+            )
+        }
+    }
+
+    @discardableResult
+    public func wait(
+        forNanoseconds duration: UInt64
+    ) async -> TaskTimeoutResult {
+        guard locked else { return .success }
+        await withTaskGroup(of: Void.self) { group in
+            group.addTask { [weak self] in await self?.wait() }
+            group.addTask {
+                do {
+                    try await Task.sleep(nanoseconds: duration)
+                } catch {}
+            }
+
+            for await _ in group.prefix(1) {
+                group.cancelAll()
+            }
+        }
+        return locked ? .timedOut : .success
+    }
+}

Sources/AsyncObject/AsyncSemaphore.swift

@@ -0,0 +1,79 @@
+import Foundation
+import OrderedCollections
+
+actor AsyncSemaphore {
+    private typealias Continuation = UnsafeContinuation<Void, Error>
+    private var continuations: OrderedDictionary<UUID, Continuation> = [:]
+    private var limit: UInt
+    private var count: Int
+
+    private func addContinuation(
+        _ continuation: Continuation,
+        withKey key: UUID
+    ) {
+        continuations[key] = continuation
+    }
+
+    private func removeContinuation(withKey key: UUID) {
+        continuations.removeValue(forKey: key)
+    }
+
+    public init(value count: UInt = 0) {
+        self.limit = count + 1
+        self.count = Int(limit)
+    }
+
+    public func signal() {
+        guard count < limit else { return }
+        count += 1
+        guard !continuations.isEmpty else { return }
+        let (_, continuation) = continuations.removeFirst()
+        continuation.resume()
+    }
+
+    public func wait() async {
+        count -= 1
+        if count > 0 { return }
+        let key = UUID()
+        do {
+            try await withUnsafeThrowingContinuationCancellationHandler(
+                handler: { (continuation: Continuation) in
+                    Task { await removeContinuation(withKey: key) }
+                },
+                { addContinuation($0, withKey: key) }
+            )
+        } catch {
+            debugPrint(
+                "Wait on semaphore for continuation task with key: \(key)"
+                + " cancelled with error \(error)"
+            )
+        }
+    }
+
+    @discardableResult
+    public func wait(
+        forNanoseconds duration: UInt64
+    ) async -> TaskTimeoutResult {
+        var timedOut = true
+        await withTaskGroup(of: Bool.self) { group in
+            group.addTask {
+                [weak self] in await self?.wait()
+                return true
+            }
+            group.addTask {
+                do {
+                    try await Task.sleep(nanoseconds: duration)
+                    return false
+                } catch {
+                    return true
+                }
+            }
+
+            for await result in group.prefix(1) {
+                timedOut = !result
+                group.cancelAll()
+            }
+        }
+        return timedOut ? .timedOut : .success
+    }
+}

Sources/AsyncObject/TaskTimeoutResult.swift

@@ -0,0 +1,4 @@
+public enum TaskTimeoutResult {
+    case success
+    case timedOut
+}

Sources/AsyncObject/UnsafeContinuationWrapper.swift

@@ -0,0 +1,33 @@
+func withUnsafeThrowingContinuationCancellationHandler<T: Sendable>(
+    handler: @Sendable (UnsafeContinuation<T, Error>) -> Void,
+    _ fn: (UnsafeContinuation<T, Error>) -> Void
+) async throws -> T {
+    typealias Continuation = UnsafeContinuation<T, Error>
+    let wrapper = Continuation.Wrapper()
+    let value = try await withTaskCancellationHandler {
+        guard let continuation = wrapper.value else { return }
+        wrapper.cancel(withError: CancellationError())
+        handler(continuation)
+    } operation: { () -> T in
+        let value = try await withUnsafeThrowingContinuation { (c: Continuation) in
+            wrapper.value = c
+            fn(c)
+        }
+        return value
+    }
+    return value
+}
+
+extension UnsafeContinuation {
+    class Wrapper {
+        var value: UnsafeContinuation?
+
+        init(value: UnsafeContinuation? = nil) {
+            self.value = value
+        }
+
+        func cancel(withError error: E) {
+            value?.resume(throwing: error)
+        }
+    }
+}

Tests/AsyncObjectTests/AsyncMutexTests.swift

@@ -0,0 +1,75 @@
+import XCTest
+@testable import AsyncObject
+
+class AsyncMutexTests: XCTestCase {
+
+    func checkWait(
+        for mutex: AsyncMutex,
+        durationInSeconds seconds: Int = 0
+    ) async throws {
+        Task.detached {
+            try await Task.sleep(nanoseconds: UInt64(5E9))
+            await mutex.release()
+        }
+        await checkExecInterval(
+            for: { await mutex.wait() },
+            durationInSeconds: seconds
+        )
+    }
+
+    func testMutexWait() async throws {
+        let mutex = AsyncMutex()
+        try await checkWait(for: mutex, durationInSeconds: 5)
+    }
+
+    func testMutexLockAndWait() async throws {
+        let mutex = AsyncMutex(lockedInitially: false)
+        await mutex.lock()
+        try await checkWait(for: mutex, durationInSeconds: 5)
+    }
+
+    func testReleasedMutexWait() async throws {
+        let mutex = AsyncMutex(lockedInitially: false)
+        try await checkWait(for: mutex)
+    }
+
+    func testMutexWaitWithTimeout() async throws {
+        let mutex = AsyncMutex()
+        var result: TaskTimeoutResult = .success
+        await checkExecInterval(
+            for: {
+                result = await mutex.wait(forNanoseconds: UInt64(4E9))
+            },
+            durationInSeconds: 4
+        )
+        XCTAssertEqual(result, .timedOut)
+    }
+
+    func testMutexWaitSuccessWithoutTimeout() async throws {
+        let mutex = AsyncMutex()
+        var result: TaskTimeoutResult = .timedOut
+        Task.detached {
+            try await Task.sleep(nanoseconds: UInt64(5E9))
+            await mutex.release()
+        }
+        await checkExecInterval(
+            for: {
+                result = await mutex.wait(forNanoseconds: UInt64(10E9))
+            },
+            durationInSeconds: 5
+        )
+        XCTAssertEqual(result, .success)
+    }
+
+    func testReleasedMutexWaitSuccessWithoutTimeout() async throws {
+        let mutex = AsyncMutex(lockedInitially: false)
+        var result: TaskTimeoutResult = .timedOut
+        await checkExecInterval(
+            for: {
+                result = await mutex.wait(forNanoseconds: UInt64(10E9))
+            },
+            durationInSeconds: 0
+        )
+        XCTAssertEqual(result, .success)
+    }
+}

Tests/AsyncObjectTests/XCTestCase.swift

@@ -0,0 +1,28 @@
+import XCTest
+
+extension XCTestCase {
+    func checkExecInterval(
+        for task: () async throws -> Void,
+        durationInSeconds seconds: Int = 0
+    ) async rethrows {
+        let time = DispatchTime.now()
+        try await task()
+        let execTime = time.distance(to: DispatchTime.now())
+        switch execTime {
+        case .seconds(let value):
+            XCTAssertEqual(seconds, value)
+        case .microseconds(let value):
+            XCTAssertEqual(seconds, value/Int(1E6))
+        case .milliseconds(let value):
+            XCTAssertEqual(seconds, value/Int(1E3))
+        case .nanoseconds(let value):
+            XCTAssertEqual(seconds, value/Int(1E9))
+        case .never: fallthrough
+        @unknown default:
+            NSException(
+                name: NSExceptionName(rawValue: "UnExpectedInterval"),
+                reason: "UnExpected time interval"
+            ).raise()
+        }
+    }
+}