swift-mutex

swhitty · Sep 8, 2024 · 7bdb47e · 7bdb47e
1 parent 10aab78
commit 7bdb47e
Show file tree

Hide file tree

Showing 7 changed files with 552 additions and 186 deletions.
diff --git a/Sources/IdentifiableContinuation.swift b/Sources/IdentifiableContinuation.swift
@@ -52,7 +52,7 @@ public func withIdentifiableContinuation<T>(
   onCancel handler: @Sendable (IdentifiableContinuation<T, Never>.ID) -> Void
 ) async -> T {
     let id = IdentifiableContinuation<T, Never>.ID()
-    let state = AllocatedLock(initialState: (isStarted: false, isCancelled: false))
+    let state = Mutex((isStarted: false, isCancelled: false))
     nonisolated(unsafe) let body = body
     return await withTaskCancellationHandler {
         await withCheckedContinuation(isolation: isolation, function: function) {
@@ -99,7 +99,7 @@ public func withIdentifiableThrowingContinuation<T>(
   onCancel handler: @Sendable (IdentifiableContinuation<T, any Error>.ID) -> Void
 ) async throws -> T {
     let id = IdentifiableContinuation<T, any Error>.ID()
-    let state = AllocatedLock(initialState: (isStarted: false, isCancelled: false))
+    let state = Mutex((isStarted: false, isCancelled: false))
     nonisolated(unsafe) let body = body
     return try await withTaskCancellationHandler {
         try await withCheckedThrowingContinuation(isolation: isolation, function: function) {
@@ -148,7 +148,7 @@ public func withIdentifiableContinuation<T>(
   onCancel handler: @Sendable (IdentifiableContinuation<T, Never>.ID) -> Void
 ) async -> T {
     let id = IdentifiableContinuation<T, Never>.ID()
-    let state = AllocatedLock(initialState: (isStarted: false, isCancelled: false))
+    let state = Mutex((isStarted: false, isCancelled: false))
     return await withTaskCancellationHandler {
         await withCheckedContinuation(function: function) {
             let continuation = IdentifiableContinuation(id: id, continuation: $0)
@@ -197,7 +197,7 @@ public func withIdentifiableThrowingContinuation<T>(
   onCancel handler: @Sendable (IdentifiableContinuation<T, any Error>.ID) -> Void
 ) async throws -> T {
     let id = IdentifiableContinuation<T, any Error>.ID()
-    let state = AllocatedLock(initialState: (isStarted: false, isCancelled: false))
+    let state = Mutex((isStarted: false, isCancelled: false))
     return try await withTaskCancellationHandler {
         try await withCheckedThrowingContinuation(function: function) {
             let continuation = IdentifiableContinuation(id: id, continuation: $0)

diff --git a/Sources/AllocatedLock.swift → Sources/Mutex.swift b/Sources/AllocatedLock.swift → Sources/Mutex.swift
@@ -1,14 +1,14 @@
 //
-//  AllocatedLock.swift
-//  AllocatedLock
+//  Mutex.swift
+//  swift-mutex
 //
 //  Created by Simon Whitty on 10/04/2023.
 //  Copyright 2023 Simon Whitty
 //
 //  Distributed under the permissive MIT license
 //  Get the latest version from here:
 //
-//  https://github.com/swhitty/AllocatedLock
+//  https://github.com/swhitty/swift-mutex
 //
 //  Permission is hereby granted, free of charge, to any person obtaining a copy
 //  of this software and associated documentation files (the "Software"), to deal
@@ -29,82 +29,99 @@
 //  SOFTWARE.
 //
 
-// Backports the Swift interface around os_unfair_lock_t available in recent Darwin platforms
-//
+// Backports the Swift 6.0 Mutex API
 @usableFromInline
-struct AllocatedLock<State>: @unchecked Sendable {
-
-    @usableFromInline
+package struct Mutex<Value>: @unchecked Sendable {
     let storage: Storage
+}
+
+#if compiler(>=6)
+package extension Mutex {
 
     @usableFromInline
-    init(initialState: State) {
-        self.storage = Storage(initialState: initialState)
+    init(_ initialValue: consuming sending Value) {
+        self.storage = Storage(initialValue)
     }
 
-    @inlinable
-    func withLock<R>(_ body: @Sendable (inout State) throws -> R) rethrows -> R where R: Sendable {
+    @usableFromInline
+    borrowing func withLock<Result, E: Error>(
+        _ body: (inout sending Value) throws(E) -> sending Result
+    ) throws(E) -> sending Result {
         storage.lock()
         defer { storage.unlock() }
-        return try body(&storage.state)
+        return try body(&storage.value)
     }
-}
-
-extension AllocatedLock where State == Void {
 
-    init() {
-        self.storage = Storage(initialState: ())
+    @usableFromInline
+    borrowing func withLockIfAvailable<Result, E>(
+        _ body: (inout sending Value) throws(E) -> sending Result
+    ) throws(E) -> sending Result? where E: Error {
+        guard storage.tryLock() else { return nil }
+        defer { storage.unlock() }
+        return try body(&storage.value)
     }
+}
+#else
+package extension Mutex {
 
-    @inlinable @available(*, noasync)
-    func lock() {
-        storage.lock()
+    @usableFromInline
+    init(_ initialValue: Value) {
+        self.storage = Storage(initialValue)
     }
 
-    @inlinable @available(*, noasync)
-    func unlock() {
-        storage.unlock()
+    @usableFromInline
+    borrowing func withLock<Result>(
+        _ body: (inout Value) throws -> Result
+    ) rethrows -> Result {
+        storage.lock()
+        defer { storage.unlock() }
+        return try body(&storage.value)
     }
 
-    @inlinable
-    func withLock<R>(_ body: @Sendable () throws -> R) rethrows -> R where R: Sendable {
-        storage.lock()
+    @usableFromInline
+    borrowing func withLockIfAvailable<Result>(
+        _ body: (inout Value) throws -> Result
+    ) rethrows -> Result? {
+        guard storage.tryLock() else { return nil }
         defer { storage.unlock() }
-        return try body()
+        return try body(&storage.value)
     }
 }
+#endif
 
 #if canImport(Darwin)
 
 import struct os.os_unfair_lock_t
 import struct os.os_unfair_lock
 import func os.os_unfair_lock_lock
 import func os.os_unfair_lock_unlock
+import func os.os_unfair_lock_trylock
+
+extension Mutex {
 
-extension AllocatedLock {
-    @usableFromInline
     final class Storage {
         private let _lock: os_unfair_lock_t
 
-        @usableFromInline
-        var state: State
+        var value: Value
 
-        init(initialState: State) {
+        init(_ initialValue: Value) {
             self._lock = .allocate(capacity: 1)
             self._lock.initialize(to: os_unfair_lock())
-            self.state = initialState
+            self.value = initialValue
         }
 
-        @usableFromInline
         func lock() {
             os_unfair_lock_lock(_lock)
         }
 
-        @usableFromInline
         func unlock() {
             os_unfair_lock_unlock(_lock)
         }
 
+        func tryLock() -> Bool {
+            os_unfair_lock_trylock(_lock)
+        }
+
         deinit {
             self._lock.deinitialize(count: 1)
             self._lock.deallocate()
@@ -116,35 +133,36 @@ extension AllocatedLock {
 
 import Glibc
 
-extension AllocatedLock {
-    @usableFromInline
+extension Mutex {
+
     final class Storage {
         private let _lock: UnsafeMutablePointer<pthread_mutex_t>
 
-        @usableFromInline
-        var state: State
+        var value: Value
 
-        init(initialState: State) {
+        init(_ initialValue: Value) {
             var attr = pthread_mutexattr_t()
             pthread_mutexattr_init(&attr)
             self._lock = .allocate(capacity: 1)
             let err = pthread_mutex_init(self._lock, &attr)
             precondition(err == 0, "pthread_mutex_init error: \(err)")
-            self.state = initialState
+            self.value = initialValue
         }
 
-        @usableFromInline
         func lock() {
             let err = pthread_mutex_lock(_lock)
             precondition(err == 0, "pthread_mutex_lock error: \(err)")
         }
 
-        @usableFromInline
         func unlock() {
             let err = pthread_mutex_unlock(_lock)
             precondition(err == 0, "pthread_mutex_unlock error: \(err)")
         }
 
+        func tryLock() -> Bool {
+            pthread_mutex_trylock(_lock) == 0
+        }
+
         deinit {
             let err = pthread_mutex_destroy(self._lock)
             precondition(err == 0, "pthread_mutex_destroy error: \(err)")
@@ -158,29 +176,30 @@ extension AllocatedLock {
 import ucrt
 import WinSDK
 
-extension AllocatedLock {
-    @usableFromInline
+extension Mutex {
+
     final class Storage {
         private let _lock: UnsafeMutablePointer<SRWLOCK>
 
-        @usableFromInline
-        var state: State
+        var value: Value
 
-        init(initialState: State) {
+        init(_ initialValue: Value) {
             self._lock = .allocate(capacity: 1)
             InitializeSRWLock(self._lock)
-            self.state = initialState
+            self.value = initialValue
         }
 
-        @usableFromInline
         func lock() {
             AcquireSRWLockExclusive(_lock)
         }
 
-        @usableFromInline
         func unlock() {
             ReleaseSRWLockExclusive(_lock)
         }
+
+        func tryLock() -> Bool {
+            TryAcquireSRWLockExclusive(_lock)
+        }
     }
 }
 

diff --git a/Tests/AllocatedLockTests.swift b/Tests/AllocatedLockTests.swift