Dont hold lock over continuation in NIOThrowingAsyncSequenceProducer (#3454)

Motivation:

The various 'withMumbleContinuation' APIs are supposed to be invoked
synchronously with the caller. This assumption allows a lock to be
acquired before the call and released from the body of the
'withMumbleContinuation' after e.g. storing the continuation. However
this isn't the case and the job may be re-enqueued on the executor
meaning that this is pattern is vulnerable to deadlocks.

Modifications:

- Drop and reacquire the lock in the NIOThrowingAsyncSequenceProducer.
- Merge 'next()' and 'next(for:)' methods in the state machine into a
single func; this reduces the amount of duplicated logic across the two
functions.

Result:

Lower chance of deadlock

Author: glbrntt

Sources/NIOCore/AsyncSequences/NIOThrowingAsyncSequenceProducer.swift

@@ -566,36 +566,33 @@ extension NIOThrowingAsyncSequenceProducer {
         @inlinable
         internal func next() async throws -> Element? {
             try await withTaskCancellationHandler { () async throws -> Element? in
-                let unsafe = self._state.unsafe
-                unsafe.lock()
-
-                let action = unsafe.withValueAssumingLockIsAcquired {
-                    $0.stateMachine.next()
+                let (action, delegate) = self._state.withLockedValue { state -> (StateMachine.NextAction, Delegate?) in
+                    let action = state.stateMachine.next(continuation: nil)
+                    switch action {
+                    case .returnElement, .returnCancellationError, .returnNil, .suspendTask_withoutContinuationOnly:
+                        return (action, nil)
+                    case .returnElementAndCallProduceMore:
+                        return (action, state.delegate)
+                    case .returnFailureAndCallDidTerminate:
+                        let delegate = state.delegate
+                        state.delegate = nil
+                        return (action, delegate)
+                    case .produceMore_withContinuationOnly, .doNothing_withContinuationOnly:
+                        // Can't be returned when 'next(continuation:)' is called with no
+                        // continuation.
+                        fatalError()
+                    }
                 }
 
                 switch action {
                 case .returnElement(let element):
-                    unsafe.unlock()
                     return element
 
                 case .returnElementAndCallProduceMore(let element):
-                    let delegate = unsafe.withValueAssumingLockIsAcquired {
-                        $0.delegate
-                    }
-                    unsafe.unlock()
-
                     delegate?.produceMore()
-
                     return element
 
                 case .returnFailureAndCallDidTerminate(let failure):
-                    let delegate = unsafe.withValueAssumingLockIsAcquired {
-                        let delegate = $0.delegate
-                        $0.delegate = nil
-                        return delegate
-                    }
-                    unsafe.unlock()
-
                     delegate?.didTerminate()
 
                     switch failure {
@@ -607,7 +604,6 @@ extension NIOThrowingAsyncSequenceProducer {
                     }
 
                 case .returnCancellationError:
-                    unsafe.unlock()
                     // We have deprecated the generic Failure type in the public API and Failure should
                     // now be `Swift.Error`. However, if users have not migrated to the new API they could
                     // still use a custom generic Error type and this cast might fail.
@@ -622,42 +618,79 @@ extension NIOThrowingAsyncSequenceProducer {
                     return nil
 
                 case .returnNil:
-                    unsafe.unlock()
                     return nil
 
-                case .suspendTask:
-                    // It is safe to hold the lock across this method
-                    // since the closure is guaranteed to be run straight away
+                case .produceMore_withContinuationOnly, .doNothing_withContinuationOnly:
+                    // Can't be returned when 'next(continuation:)' is called with no
+                    // continuation.
+                    fatalError()
+
+                case .suspendTask_withoutContinuationOnly:
+                    // Holding the lock here *should* be safe but because of a bug in the runtime
+                    // it isn't, so drop the lock, create the continuation and then try again.
+                    //
+                    // See https://github.com/swiftlang/swift/issues/85668
                     return try await withCheckedThrowingContinuation {
                         (continuation: CheckedContinuation<Element?, any Error>) in
-                        let (action, callDidSuspend) = unsafe.withValueAssumingLockIsAcquired {
-                            let action = $0.stateMachine.next(for: continuation)
-                            let callDidSuspend = $0.didSuspend != nil
-                            return (action, callDidSuspend)
+                        let (action, delegate, didSuspend) = self._state.withLockedValue { state in
+                            let action = state.stateMachine.next(continuation: continuation)
+                            let delegate = state.delegate
+                            return (action, delegate, state.didSuspend)
                         }
 
                         switch action {
-                        case .callProduceMore:
-                            let delegate = unsafe.withValueAssumingLockIsAcquired {
-                                $0.delegate
+                        case .returnElement(let element):
+                            continuation.resume(returning: element)
+
+                        case .returnElementAndCallProduceMore(let element):
+                            delegate?.produceMore()
+                            continuation.resume(returning: element)
+
+                        case .returnFailureAndCallDidTerminate(let failure):
+                            delegate?.didTerminate()
+                            switch failure {
+                            case .some(let error):
+                                continuation.resume(throwing: error)
+                            case .none:
+                                continuation.resume(returning: nil)
                             }
-                            unsafe.unlock()
 
+                        case .returnCancellationError:
+                            // We have deprecated the generic Failure type in the public API and Failure should
+                            // now be `Swift.Error`. However, if users have not migrated to the new API they could
+                            // still use a custom generic Error type and this cast might fail.
+                            // In addition, we use `NIOThrowingAsyncSequenceProducer` in the implementation of the
+                            // non-throwing variant `NIOAsyncSequenceProducer` where `Failure` will be `Never` and
+                            // this cast will fail as well.
+                            // Everything is marked @inlinable and the Failure type is known at compile time,
+                            // therefore this cast should be optimised away in release build.
+                            if let error = CancellationError() as? Failure {
+                                continuation.resume(throwing: error)
+                            } else {
+                                continuation.resume(returning: nil)
+                            }
+
+                        case .returnNil:
+                            continuation.resume(returning: nil)
+
+                        case .produceMore_withContinuationOnly:
                             delegate?.produceMore()
 
-                        case .none:
-                            unsafe.unlock()
-                        }
+                        case .doNothing_withContinuationOnly:
+                            ()
 
-                        if callDidSuspend {
-                            let didSuspend = self._state.withLockedValue { $0.didSuspend }
-                            didSuspend?()
+                        case .suspendTask_withoutContinuationOnly:
+                            // Can't be returned when 'next(continuation:)' is called with a
+                            // continuation.
+                            fatalError()
                         }
+
+                        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 unlock
                 // in cases where we race with a cancellation handler
                 let (delegate, action): (Delegate?, NIOThrowingAsyncSequenceProducer.StateMachine.CancelledAction) =
                     self._state.withLockedValue {
@@ -1167,14 +1200,25 @@ extension NIOThrowingAsyncSequenceProducer {
             case returnCancellationError
             /// Indicates that the `nil` should be returned to the caller.
             case returnNil
-            /// Indicates that the `Task` of the caller should be suspended.
-            case suspendTask
+
+            /// Indicates that the `Task` of the caller should be suspended. Only returned when
+            /// `next(continuation:)` is called with a non-nil continuation.
+            case suspendTask_withoutContinuationOnly
+
+            /// Indicates that caller should produce more values. Only returned when
+            /// `next(continuation:)` is called with `nil`.
+            case produceMore_withContinuationOnly
+            /// Indicates that caller shouldn't do anything. Only returned
+            /// when `next(continuation:)` is called with `nil`.
+            case doNothing_withContinuationOnly
         }
 
         @inlinable
-        mutating func next() -> NextAction {
+        mutating func next(continuation: CheckedContinuation<Element?, Error>?) -> NextAction {
             switch self._state {
             case .initial(let backPressureStrategy, let iteratorInitialized):
+                precondition(continuation == nil)
+
                 // We are not interacting with the back-pressure strategy here because
                 // we are doing this inside `next(:)`
                 self._state = .streaming(
@@ -1185,7 +1229,7 @@ extension NIOThrowingAsyncSequenceProducer {
                     iteratorInitialized: iteratorInitialized
                 )
 
-                return .suspendTask
+                return .suspendTask_withoutContinuationOnly
 
             case .streaming(_, _, .some, _, _):
                 // We have multiple AsyncIterators iterating the sequence
@@ -1202,7 +1246,6 @@ extension NIOThrowingAsyncSequenceProducer {
 
                 if let element = buffer.popFirst() {
                     // We have an element to fulfil the demand right away.
-
                     let shouldProduceMore = backPressureStrategy.didConsume(bufferDepth: buffer.count)
 
                     self._state = .streaming(
@@ -1220,10 +1263,27 @@ extension NIOThrowingAsyncSequenceProducer {
                         // We don't have any new demand, so we can just return the element.
                         return .returnElement(element)
                     }
+                } else if let continuation = continuation {
+                    let shouldProduceMore = backPressureStrategy.didConsume(bufferDepth: buffer.count)
+                    self._state = .streaming(
+                        backPressureStrategy: backPressureStrategy,
+                        buffer: buffer,
+                        continuation: continuation,
+                        hasOutstandingDemand: shouldProduceMore,
+                        iteratorInitialized: iteratorInitialized
+                    )
+
+                    if shouldProduceMore && !hasOutstandingDemand {
+                        return .produceMore_withContinuationOnly
+                    } else {
+                        return .doNothing_withContinuationOnly
+                    }
                 } else {
-                    // There is nothing in the buffer to fulfil the demand so we need to suspend.
-                    // We are not interacting with the back-pressure strategy here because
-                    // we are doing this inside `next(:)`
+                    // This function is first called without a continuation (i.e. this branch), if
+                    // the buffer is empty the caller needs to re-call this function with a
+                    // continuation (the branch above) so defer checking the backpressure strategy
+                    // until then to minimise unnecessary delegate calls.
+
                     self._state = .streaming(
                         backPressureStrategy: backPressureStrategy,
                         buffer: buffer,
@@ -1232,7 +1292,7 @@ extension NIOThrowingAsyncSequenceProducer {
                         iteratorInitialized: iteratorInitialized
                     )
 
-                    return .suspendTask
+                    return .suspendTask_withoutContinuationOnly
                 }
 
             case .sourceFinished(var buffer, let iteratorInitialized, let failure):
@@ -1265,56 +1325,6 @@ extension NIOThrowingAsyncSequenceProducer {
                 preconditionFailure("Invalid state")
             }
         }
-
-        /// Actions returned by `next(for:)`.
-        @usableFromInline
-        enum NextForContinuationAction: Sendable {
-            /// Indicates that ``NIOAsyncSequenceProducerDelegate/produceMore()`` should be called.
-            case callProduceMore
-            /// Indicates that nothing should be done.
-            case none
-        }
-
-        @inlinable
-        mutating func next(for continuation: CheckedContinuation<Element?, Error>) -> NextForContinuationAction {
-            switch self._state {
-            case .initial:
-                // We are transitioning away from the initial state in `next()`
-                preconditionFailure("Invalid state")
-
-            case .streaming(
-                var backPressureStrategy,
-                let buffer,
-                .none,
-                let hasOutstandingDemand,
-                let iteratorInitialized
-            ):
-                precondition(buffer.isEmpty, "Expected an empty buffer")
-
-                self._state = .modifying
-                let shouldProduceMore = backPressureStrategy.didConsume(bufferDepth: buffer.count)
-
-                self._state = .streaming(
-                    backPressureStrategy: backPressureStrategy,
-                    buffer: buffer,
-                    continuation: continuation,
-                    hasOutstandingDemand: shouldProduceMore,
-                    iteratorInitialized: iteratorInitialized
-                )
-
-                if shouldProduceMore && !hasOutstandingDemand {
-                    return .callProduceMore
-                } else {
-                    return .none
-                }
-
-            case .streaming(_, _, .some(_), _, _), .sourceFinished, .finished, .cancelled:
-                preconditionFailure("This should have already been handled by `next()`")
-
-            case .modifying:
-                preconditionFailure("Invalid state")
-            }
-        }
     }
 }
 

docs/workarounds.md

@@ -17,3 +17,4 @@ or cases where the Swift compiler was unable to sufficiently optimize the code:
 - https://github.com/apple/swift-nio/pull/2046
 - https://github.com/apple/swift-nio/pull/3303
 - https://github.com/apple/swift-nio/pull/3452
+- https://github.com/apple/swift-nio/pull/3454