Adopt NIOLock

Sources/HTTPTypes/HTTPFields.swift

@@ -12,16 +12,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#if canImport(os.lock)
-import os.lock
-#elseif canImport(Glibc)
-import Glibc
-#elseif canImport(Musl)
-import Musl
-#elseif canImport(WASILibc)
-import WASILibc
-#endif
-
 /// A collection of HTTP fields. It is used in `HTTPRequest` and `HTTPResponse`, and can also be
 /// used as HTTP trailer fields.
 ///
@@ -34,77 +24,33 @@ public struct HTTPFields: Sendable, Hashable {
     private final class _Storage: @unchecked Sendable, Hashable {
         var fields: [(field: HTTPField, next: UInt16)] = []
         var index: [String: (first: UInt16, last: UInt16)]? = [:]
-        #if canImport(os.lock)
-        let lock = UnsafeMutablePointer<os_unfair_lock>.allocate(capacity: 1)
-        #else
-        let lock = UnsafeMutablePointer<pthread_mutex_t>.allocate(capacity: 1)
-        #endif
-
-        init() {
-            #if canImport(os.lock)
-            self.lock.initialize(to: os_unfair_lock())
-            #else
-            let err = pthread_mutex_init(self.lock, nil)
-            precondition(err == 0, "pthread_mutex_init failed with error \(err)")
-            #endif
-        }
-
-        deinit {
-            #if !canImport(os.lock)
-            let err = pthread_mutex_destroy(self.lock)
-            precondition(err == 0, "pthread_mutex_destroy failed with error \(err)")
-            #endif
-            self.lock.deallocate()
-        }
+        let lock = LockStorage.create(value: ())
 
         var ensureIndex: [String: (first: UInt16, last: UInt16)] {
-            #if canImport(os.lock)
-            os_unfair_lock_lock(self.lock)
-            defer { os_unfair_lock_unlock(self.lock) }
-            #else
-            let err = pthread_mutex_lock(self.lock)
-            precondition(err == 0, "pthread_mutex_lock failed with error \(err)")
-            defer {
-                let err = pthread_mutex_unlock(self.lock)
-                precondition(err == 0, "pthread_mutex_unlock failed with error \(err)")
-            }
-            #endif
-            if let index = self.index {
-                return index
-            }
-            var newIndex = [String: (first: UInt16, last: UInt16)]()
-            for index in self.fields.indices {
-                let name = self.fields[index].field.name.canonicalName
-                self.fields[index].next = .max
-                if let lastIndex = newIndex[name]?.last {
-                    self.fields[Int(lastIndex)].next = UInt16(index)
+            self.lock.withLockedValue { _ in
+                if let index = self.index {
+                    return index
                 }
-                newIndex[name, default: (first: UInt16(index), last: 0)].last = UInt16(index)
+                var newIndex = [String: (first: UInt16, last: UInt16)]()
+                for index in self.fields.indices {
+                    let name = self.fields[index].field.name.canonicalName
+                    self.fields[index].next = .max
+                    if let lastIndex = newIndex[name]?.last {
+                        self.fields[Int(lastIndex)].next = UInt16(index)
+                    }
+                    newIndex[name, default: (first: UInt16(index), last: 0)].last = UInt16(index)
+                }
+                self.index = newIndex
+                return newIndex
             }
-            self.index = newIndex
-            return newIndex
         }
 
         func copy() -> _Storage {
             let newStorage = _Storage()
             newStorage.fields = self.fields
-            #if canImport(os.lock)
-            os_unfair_lock_lock(self.lock)
-            #else
-            do {
-                let err = pthread_mutex_lock(self.lock)
-                precondition(err == 0, "pthread_mutex_lock failed with error \(err)")
-            }
-            #endif
-            newStorage.index = self.index
-            #if canImport(os.lock)
-            os_unfair_lock_unlock(self.lock)
-            #else
-            do {
-                let err = pthread_mutex_unlock(self.lock)
-                precondition(err == 0, "pthread_mutex_unlock failed with error \(err)")
+            self.lock.withLockedValue { _ in
+                newStorage.index = self.index
             }
-            #endif
             return newStorage
         }
 

Sources/HTTPTypes/NIOLock.swift

@@ -0,0 +1,271 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2017-2024 Apple Inc. and the SwiftNIO project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+#if canImport(Darwin)
+import Darwin
+#elseif os(Windows)
+import ucrt
+import WinSDK
+#elseif canImport(Glibc)
+import Glibc
+#elseif canImport(Musl)
+import Musl
+#elseif canImport(Bionic)
+import Bionic
+#elseif canImport(WASILibc)
+import WASILibc
+#if canImport(wasi_pthread)
+import wasi_pthread
+#endif
+#else
+#error("The concurrency NIOLock module was unable to identify your C library.")
+#endif
+
+#if os(Windows)
+@usableFromInline
+typealias LockPrimitive = SRWLOCK
+#elseif canImport(Darwin)
+@usableFromInline
+typealias LockPrimitive = os_unfair_lock
+#else
+@usableFromInline
+typealias LockPrimitive = pthread_mutex_t
+#endif
+
+@usableFromInline
+enum LockOperations {}
+
+extension LockOperations {
+    @inlinable
+    static func create(_ mutex: UnsafeMutablePointer<LockPrimitive>) {
+        mutex.assertValidAlignment()
+
+        #if os(Windows)
+        InitializeSRWLock(mutex)
+        #elseif canImport(Darwin)
+        mutex.initialize(to: .init())
+        #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded))
+        var attr = pthread_mutexattr_t()
+        pthread_mutexattr_init(&attr)
+
+        let err = pthread_mutex_init(mutex, &attr)
+        precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)")
+        #endif
+    }
+
+    @inlinable
+    static func destroy(_ mutex: UnsafeMutablePointer<LockPrimitive>) {
+        mutex.assertValidAlignment()
+
+        #if os(Windows)
+        // SRWLOCK does not need to be free'd
+        #elseif canImport(Darwin)
+        // os_unfair_lock does not need to be free'd
+        #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded))
+        let err = pthread_mutex_destroy(mutex)
+        precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)")
+        #endif
+    }
+
+    @inlinable
+    static func lock(_ mutex: UnsafeMutablePointer<LockPrimitive>) {
+        mutex.assertValidAlignment()
+
+        #if os(Windows)
+        AcquireSRWLockExclusive(mutex)
+        #elseif canImport(Darwin)
+        os_unfair_lock_lock(mutex)
+        #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded))
+        let err = pthread_mutex_lock(mutex)
+        precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)")
+        #endif
+    }
+
+    @inlinable
+    static func unlock(_ mutex: UnsafeMutablePointer<LockPrimitive>) {
+        mutex.assertValidAlignment()
+
+        #if os(Windows)
+        ReleaseSRWLockExclusive(mutex)
+        #elseif canImport(Darwin)
+        os_unfair_lock_unlock(mutex)
+        #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded))
+        let err = pthread_mutex_unlock(mutex)
+        precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)")
+        #endif
+    }
+}
+
+// Tail allocate both the mutex and a generic value using ManagedBuffer.
+// Both the header pointer and the elements pointer are stable for
+// the class's entire lifetime.
+//
+// However, for safety reasons, we elect to place the lock in the "elements"
+// section of the buffer instead of the head. The reasoning here is subtle,
+// so buckle in.
+//
+// _As a practical matter_, the implementation of ManagedBuffer ensures that
+// the pointer to the header is stable across the lifetime of the class, and so
+// each time you call `withUnsafeMutablePointers` or `withUnsafeMutablePointerToHeader`
+// the value of the header pointer will be the same. This is because ManagedBuffer uses
+// `Builtin.addressOf` to load the value of the header, and that does ~magic~ to ensure
+// that it does not invoke any weird Swift accessors that might copy the value.
+//
+// _However_, the header is also available via the `.header` field on the ManagedBuffer.
+// This presents a problem! The reason there's an issue is that `Builtin.addressOf` and friends
+// do not interact with Swift's exclusivity model. That is, the various `with` functions do not
+// conceptually trigger a mutating access to `.header`. For elements this isn't a concern because
+// there's literally no other way to perform the access, but for `.header` it's entirely possible
+// to accidentally recursively read it.
+//
+// Our implementation is free from these issues, so we don't _really_ need to worry about it.
+// However, out of an abundance of caution, we store the Value in the header, and the LockPrimitive
+// in the trailing elements. We still don't use `.header`, but it's better to be safe than sorry,
+// and future maintainers will be happier that we were cautious.
+//
+// See also: https://github.com/apple/swift/pull/40000
+@usableFromInline
+final class LockStorage<Value>: ManagedBuffer<Value, LockPrimitive> {
+
+    @inlinable
+    static func create(value: Value) -> Self {
+        let buffer = Self.create(minimumCapacity: 1) { _ in
+            value
+        }
+        // Intentionally using a force cast here to avoid a miss compiliation in 5.10.
+        // This is as fast as an unsafeDownCast since ManagedBuffer is inlined and the optimizer
+        // can eliminate the upcast/downcast pair
+        let storage = buffer as! Self
+
+        storage.withUnsafeMutablePointers { _, lockPtr in
+            LockOperations.create(lockPtr)
+        }
+
+        return storage
+    }
+
+    @inlinable
+    func lock() {
+        self.withUnsafeMutablePointerToElements { lockPtr in
+            LockOperations.lock(lockPtr)
+        }
+    }
+
+    @inlinable
+    func unlock() {
+        self.withUnsafeMutablePointerToElements { lockPtr in
+            LockOperations.unlock(lockPtr)
+        }
+    }
+
+    @inlinable
+    deinit {
+        self.withUnsafeMutablePointerToElements { lockPtr in
+            LockOperations.destroy(lockPtr)
+        }
+    }
+
+    @inlinable
+    func withLockPrimitive<T>(_ body: (UnsafeMutablePointer<LockPrimitive>) throws -> T) rethrows -> T {
+        try self.withUnsafeMutablePointerToElements { lockPtr in
+            try body(lockPtr)
+        }
+    }
+
+    @inlinable
+    func withLockedValue<T>(_ mutate: (inout Value) throws -> T) rethrows -> T {
+        try self.withUnsafeMutablePointers { valuePtr, lockPtr in
+            LockOperations.lock(lockPtr)
+            defer { LockOperations.unlock(lockPtr) }
+            return try mutate(&valuePtr.pointee)
+        }
+    }
+}
+
+extension LockStorage: @unchecked Sendable {}
+
+/// A threading lock based on `libpthread` instead of `libdispatch`.
+///
+/// - note: ``NIOLock`` has reference semantics.
+///
+/// This object provides a lock on top of a single `pthread_mutex_t`. This kind
+/// of lock is safe to use with `libpthread`-based threading models, such as the
+/// one used by NIO. On Windows, the lock is based on the substantially similar
+/// `SRWLOCK` type.
+public struct NIOLock {
+    @usableFromInline
+    internal let _storage: LockStorage<Void>
+
+    /// Create a new lock.
+    @inlinable
+    public init() {
+        self._storage = .create(value: ())
+    }
+
+    /// Acquire the lock.
+    ///
+    /// Whenever possible, consider using `withLock` instead of this method and
+    /// `unlock`, to simplify lock handling.
+    @inlinable
+    public func lock() {
+        self._storage.lock()
+    }
+
+    /// Release the lock.
+    ///
+    /// Whenever possible, consider using `withLock` instead of this method and
+    /// `lock`, to simplify lock handling.
+    @inlinable
+    public func unlock() {
+        self._storage.unlock()
+    }
+
+    @inlinable
+    internal func withLockPrimitive<T>(_ body: (UnsafeMutablePointer<LockPrimitive>) throws -> T) rethrows -> T {
+        try self._storage.withLockPrimitive(body)
+    }
+}
+
+extension NIOLock {
+    /// Acquire the lock for the duration of the given block.
+    ///
+    /// This convenience method should be preferred to `lock` and `unlock` in
+    /// most situations, as it ensures that the lock will be released regardless
+    /// of how `body` exits.
+    ///
+    /// - Parameter body: The block to execute while holding the lock.
+    /// - Returns: The value returned by the block.
+    @inlinable
+    public func withLock<T>(_ body: () throws -> T) rethrows -> T {
+        self.lock()
+        defer {
+            self.unlock()
+        }
+        return try body()
+    }
+
+    @inlinable
+    public func withLockVoid(_ body: () throws -> Void) rethrows {
+        try self.withLock(body)
+    }
+}
+
+extension NIOLock: Sendable {}
+
+extension UnsafeMutablePointer {
+    @inlinable
+    func assertValidAlignment() {
+        assert(UInt(bitPattern: self) % UInt(MemoryLayout<Pointee>.alignment) == 0)
+    }
+}