Add manual control to NIOLockedValueBox

Motivation:

NIOLockedValueBox has a 'safer' API than NIOLock as it only provides
scoped access to its boxed value. NIOLock requires users to only access
protected state while the lock is acquired. As such NIOLockedValueBox
should be preferred where possible. However, there are cases where
manual control must be used (such as storing a continuation) and users
must use a NIOLock for this.

There are two downsides to this:

1. All other access to the protected state must use the NIOLock API
   putting the onus on the developer to only access the protected state
   while the lock is held.
2. NIOLock can't store its protected state inline which typically
   results in users storing it on a class.

Modifications:

- Add an 'unsafe' view to NIOLockedValueBox which allows users to
  manually control the lock and access its protected state
- Update NIOAsyncWriter and NIOThrowingAsyncSequenceProducer to use NIOLockedValueBox

Result:

- Safer locking API is used in more places
- Fewer allocations

Author: glbrntt

Benchmarks/Thresholds/5.10/NIOCoreBenchmarks.NIOAsyncChannel.init.p90.json

@@ -1,3 +1,3 @@
 {
-  "mallocCountTotal" : 10
+  "mallocCountTotal" : 8
 }

Benchmarks/Thresholds/5.8/NIOCoreBenchmarks.NIOAsyncChannel.init.p90.json

@@ -1,3 +1,3 @@
 {
-  "mallocCountTotal" : 10
+  "mallocCountTotal" : 8
 }

Benchmarks/Thresholds/5.9/NIOCoreBenchmarks.NIOAsyncChannel.init.p90.json

@@ -1,3 +1,3 @@
 {
-  "mallocCountTotal" : 10
+  "mallocCountTotal" : 8
 }

Benchmarks/Thresholds/nightly-main/NIOCoreBenchmarks.NIOAsyncChannel.init.p90.json

@@ -1,3 +1,3 @@
 {
-  "mallocCountTotal" : 10
+  "mallocCountTotal" : 8
 }

Benchmarks/Thresholds/nightly-next/NIOCoreBenchmarks.NIOAsyncChannel.init.p90.json

@@ -1,3 +1,3 @@
 {
-  "mallocCountTotal" : 10
+  "mallocCountTotal" : 8
 }

Sources/NIOConcurrencyHelpers/NIOLockedValueBox.swift

@@ -37,6 +37,43 @@ public struct NIOLockedValueBox<Value> {
     public func withLockedValue<T>(_ mutate: (inout Value) throws -> T) rethrows -> T {
         return try self._storage.withLockedValue(mutate)
     }
+
+    /// Provides an unsafe view over the lock and its value.
+    ///
+    /// This can be beneficial when you require fine grained control over the lock in some
+    /// situations but don't want lose the benefits of ``withLockedValue(_:)`` in others by
+    /// switching to ``NIOLock``.
+    public var unsafe: Unsafe {
+        Unsafe(_storage: self._storage)
+    }
+
+    /// Provides an unsafe view over the lock and its value.
+    public struct Unsafe {
+        @usableFromInline
+        let _storage: LockStorage<Value>
+
+        /// Manually acquire the lock.
+        @inlinable
+        public func lock() {
+            self._storage.lock()
+        }
+
+        /// Manually release the lock.
+        @inlinable
+        public func unlock() {
+            self._storage.unlock()
+        }
+
+        /// Mutate the value, assuming the lock has been acquired manually.
+        @inlinable
+        public func withValueAssumingLockIsAcquired<T>(
+            _ mutate: (inout Value) throws -> T
+        ) rethrows -> T {
+            return try self._storage.withUnsafeMutablePointerToHeader { value in
+                try mutate(&value.pointee)
+            }
+        }
+    }
 }
 
 extension NIOLockedValueBox: Sendable where Value: Sendable {}

Sources/NIOCore/AsyncSequences/NIOAsyncWriter.swift

@@ -414,12 +414,12 @@ extension NIOAsyncWriter {
 extension NIOAsyncWriter {
     /// This is the underlying storage of the writer. The goal of this is to synchronize the access to all state.
     @usableFromInline
-    /* fileprivate */ internal final class Storage: @unchecked Sendable {
+    /* fileprivate */ internal struct Storage: Sendable {
         /// Internal type to generate unique yield IDs.
         ///
         /// This type has reference semantics.
         @usableFromInline
-        struct YieldIDGenerator {
+        struct YieldIDGenerator: Sendable {
             /// A struct representing a unique yield ID.
             @usableFromInline
             struct YieldID: Equatable, Sendable {
@@ -445,47 +445,61 @@ extension NIOAsyncWriter {
             }
         }
 
-        /// The lock that protects our state.
-        @usableFromInline
-        /* private */ internal let _lock = NIOLock()
         /// The counter used to assign an ID to all our yields.
         @usableFromInline
         /* private */ internal let _yieldIDGenerator = YieldIDGenerator()
         /// The state machine.
         @usableFromInline
-        /* private */ internal var _stateMachine: StateMachine
+        /* private */ internal let _state: NIOLockedValueBox<State>
+
+        @usableFromInline
+        struct State: Sendable {
+            @usableFromInline
+            var stateMachine: StateMachine
+            @usableFromInline
+            var didSuspend: (@Sendable () -> Void)?
+
+            @inlinable
+            init(stateMachine: StateMachine) {
+                self.stateMachine = stateMachine
+                self.didSuspend = nil
+            }
+        }
+
         /// Hook used in testing.
         @usableFromInline
-        internal var _didSuspend: (() -> Void)?
+        internal func _setDidSuspend(_ didSuspend: (@Sendable () -> Void)?) {
+            self._state.withLockedValue {
+                $0.didSuspend = didSuspend
+            }
+        }
 
         @inlinable
         internal var isWriterFinished: Bool {
-            self._lock.withLock { self._stateMachine.isWriterFinished }
+            self._state.withLockedValue { $0.stateMachine.isWriterFinished }
         }
 
         @inlinable
         internal var isSinkFinished: Bool {
-            self._lock.withLock { self._stateMachine.isSinkFinished }
+            self._state.withLockedValue { $0.stateMachine.isSinkFinished }
         }
 
         @inlinable
         /* fileprivate */ internal init(
             isWritable: Bool,
             delegate: Delegate
         ) {
-            self._stateMachine = .init(
-                isWritable: isWritable,
-                delegate: delegate
-            )
+            let state = State(stateMachine: StateMachine(isWritable: isWritable, delegate: delegate))
+            self._state = NIOLockedValueBox(state)
         }
 
         @inlinable
         /* fileprivate */ internal func setWritability(to writability: Bool) {
             // We must not resume the continuation while holding the lock
             // because it can deadlock in combination with the underlying ulock
             // in cases where we race with a cancellation handler
-            let action = self._lock.withLock {
-                self._stateMachine.setWritability(to: writability)
+            let action = self._state.withLockedValue {
+                $0.stateMachine.setWritability(to: writability)
             }
 
             switch action {
@@ -516,39 +530,42 @@ extension NIOAsyncWriter {
 
             return try await withTaskCancellationHandler {
                 // We are manually locking here to hold the lock across the withCheckedContinuation call
-                self._lock.lock()
+                let unsafe = self._state.unsafe
+                unsafe.lock()
 
-                let action = self._stateMachine.yield(yieldID: yieldID)
+                let action = unsafe.withValueAssumingLockIsAcquired {
+                    $0.stateMachine.yield(yieldID: yieldID)
+                }
 
                 switch action {
                 case .callDidYield(let delegate):
                     // We are allocating a new Deque for every write here
-                    self._lock.unlock()
+                    unsafe.unlock()
                     delegate.didYield(contentsOf: Deque(sequence))
                     self.unbufferQueuedEvents()
                     return .yielded
 
                 case .throwError(let error):
-                    self._lock.unlock()
+                    unsafe.unlock()
                     throw error
 
                 case .suspendTask:
                     return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<StateMachine.YieldResult, Error>) in
-                        self._stateMachine.yield(
-                            continuation: continuation,
-                            yieldID: yieldID
-                        )
+                        let didSuspend = unsafe.withValueAssumingLockIsAcquired {
+                            $0.stateMachine.yield(continuation: continuation, yieldID: yieldID)
+                            return $0.didSuspend
+                        }
 
-                        self._lock.unlock()
-                        self._didSuspend?()
+                        unsafe.unlock()
+                        didSuspend?()
                     }
                 }
             } onCancel: {
                 // We must not resume the continuation while holding the lock
                 // because it can deadlock in combination with the underlying ulock
                 // in cases where we race with a cancellation handler
-                let action = self._lock.withLock {
-                    self._stateMachine.cancel(yieldID: yieldID)
+                let action = self._state.withLockedValue {
+                    $0.stateMachine.cancel(yieldID: yieldID)
                 }
 
                 switch action {
@@ -580,39 +597,41 @@ extension NIOAsyncWriter {
 
             return try await withTaskCancellationHandler {
                 // We are manually locking here to hold the lock across the withCheckedContinuation call
-                self._lock.lock()
+                let unsafe = self._state.unsafe
+                unsafe.lock()
 
-                let action = self._stateMachine.yield(yieldID: yieldID)
+                let action = unsafe.withValueAssumingLockIsAcquired {
+                    $0.stateMachine.yield(yieldID: yieldID)
+                }
 
                 switch action {
                 case .callDidYield(let delegate):
                     // We are allocating a new Deque for every write here
-                    self._lock.unlock()
+                    unsafe.unlock()
                     delegate.didYield(element)
                     self.unbufferQueuedEvents()
                     return .yielded
 
                 case .throwError(let error):
-                    self._lock.unlock()
+                    unsafe.unlock()
                     throw error
 
                 case .suspendTask:
                     return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<StateMachine.YieldResult, Error>) in
-                        self._stateMachine.yield(
-                            continuation: continuation,
-                            yieldID: yieldID
-                        )
-
-                        self._lock.unlock()
-                        self._didSuspend?()
+                        let didSuspend = unsafe.withValueAssumingLockIsAcquired {
+                            $0.stateMachine.yield(continuation: continuation, yieldID: yieldID)
+                            return $0.didSuspend
+                        }
+                        unsafe.unlock()
+                        didSuspend?()
                     }
                 }
             } onCancel: {
                 // We must not resume the continuation while holding the lock
-                // because it can deadlock in combination with the underlying ulock
+                // because it can deadlock in combination with the underlying lock
                 // in cases where we race with a cancellation handler
-                let action = self._lock.withLock {
-                    self._stateMachine.cancel(yieldID: yieldID)
+                let action = self._state.withLockedValue {
+                    $0.stateMachine.cancel(yieldID: yieldID)
                 }
 
                 switch action {
@@ -630,8 +649,8 @@ extension NIOAsyncWriter {
             // We must not resume the continuation while holding the lock
             // because it can deadlock in combination with the underlying ulock
             // in cases where we race with a cancellation handler
-            let action = self._lock.withLock {
-                self._stateMachine.writerFinish(error: error)
+            let action = self._state.withLockedValue {
+                $0.stateMachine.writerFinish(error: error)
             }
 
             switch action {
@@ -651,8 +670,8 @@ extension NIOAsyncWriter {
             // We must not resume the continuation while holding the lock
             // because it can deadlock in combination with the underlying ulock
             // in cases where we race with a cancellation handler
-            let action = self._lock.withLock {
-                self._stateMachine.sinkFinish(error: error)
+            let action = self._state.withLockedValue {
+                $0.stateMachine.sinkFinish(error: error)
             }
 
             switch action {
@@ -667,7 +686,7 @@ extension NIOAsyncWriter {
 
         @inlinable
         /* fileprivate */ internal func unbufferQueuedEvents() {
-            while let action = self._lock.withLock({ self._stateMachine.unbufferQueuedEvents()}) {
+            while let action = self._state.withLockedValue({ $0.stateMachine.unbufferQueuedEvents()}) {
                 switch action {
                 case .callDidTerminate(let delegate, let error):
                     delegate.didTerminate(error: error)
@@ -684,12 +703,12 @@ extension NIOAsyncWriter {
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
 extension NIOAsyncWriter {
     @usableFromInline
-    /* private */ internal struct StateMachine {
+    /* private */ internal struct StateMachine: Sendable {
         @usableFromInline
         typealias YieldID = Storage.YieldIDGenerator.YieldID
         /// This is a small helper struct to encapsulate the two different values for a suspended yield.
         @usableFromInline
-        /* private */ internal struct SuspendedYield {
+        /* private */ internal struct SuspendedYield: Sendable {
             /// The yield's ID.
             @usableFromInline
             var yieldID: YieldID
@@ -715,7 +734,7 @@ extension NIOAsyncWriter {
 
         /// The current state of our ``NIOAsyncWriter``.
         @usableFromInline
-        /* private */ internal enum State: CustomStringConvertible {
+        /* private */ internal enum State: Sendable, CustomStringConvertible {
             /// The initial state before either a call to ``NIOAsyncWriter/yield(contentsOf:)`` or
             /// ``NIOAsyncWriter/finish(completion:)`` happened.
             case initial(