Apple/swift-nio
1
0
Fork 0
mirror of https://github.com/apple/swift-nio.git synced 2026-05-20 20:30:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Drop and reacquire lock over continuation call (#3452)

Browse Source 
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 NIOAsyncWriter

Result:

Lower chance of deadlock
This commit is contained in:
George Barnett
2025-11-24 11:13:50 +00:00
committed by GitHub
parent 3eea09220e
commit 5bf841dde5
2 changed files with 86 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/NIOCore/AsyncSequences/NIOAsyncWriter.swift
+85 -25
View File
@@ -540,36 +540,66 @@ extension NIOAsyncWriter {
let yieldID = yieldID ?? self._yieldIDGenerator.generateUniqueYieldID()
return try await withTaskCancellationHandler {
// We are manually locking here to hold the lock across the withCheckedContinuation call
let unsafe = self._state.unsafe
unsafe.lock()
let action = unsafe.withValueAssumingLockIsAcquired {
let action = self._state.withLockedValue {
$0.stateMachine.yield(yieldID: yieldID)
}
switch action {
case .callDidYield(let delegate):
// We are allocating a new Deque for every write here
unsafe.unlock()
delegate.didYield(contentsOf: Deque(sequence))
self.unbufferQueuedEvents()
return .yielded
case .throwError(let error):
unsafe.unlock()
throw error
case .suspendTask:
// 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
//
// Dropping and reacquiring the lock may result in yields being reordered but
// only from the perspective of when this function was entered. For example:
//
// - T1 calls _yield
// - T2 calls _yield
// - T2 returns from _yield
// - T1 returns from _yield
//
// This is fine: the async writer doesn't offer any ordering guarantees for
// calls made from different threads.
//
// Within a thread there is no possibility of re-ordering as the call only
// returns once the write has been yielded.
return try await withCheckedThrowingContinuation {
(continuation: CheckedContinuation<StateMachine.YieldResult, Error>) in
let didSuspend = unsafe.withValueAssumingLockIsAcquired {
$0.stateMachine.yield(continuation: continuation, yieldID: yieldID)
return $0.didSuspend
let (action, didSuspend) = self._state.withLockedValue {
state -> (NIOAsyncWriter.StateMachine.YieldAction, (@Sendable () -> Void)?) in
let yieldAction = state.stateMachine.yield(yieldID: yieldID)
switch yieldAction {
case .callDidYield, .throwError:
return (yieldAction, nil)
case .suspendTask:
state.stateMachine.yield(continuation: continuation, yieldID: yieldID)
let didSuspend = state.didSuspend
return (yieldAction, didSuspend)
}
}
unsafe.unlock()
didSuspend?()
switch action {
case .callDidYield(let delegate):
delegate.didYield(contentsOf: Deque(sequence))
self.unbufferQueuedEvents()
continuation.resume(returning: .yielded)
case .throwError(let error):
continuation.resume(throwing: error)
case .suspendTask:
didSuspend?()
}
}
}
} onCancel: {
@@ -611,35 +641,65 @@ extension NIOAsyncWriter {
let yieldID = yieldID ?? self._yieldIDGenerator.generateUniqueYieldID()
return try await withTaskCancellationHandler {
// We are manually locking here to hold the lock across the withCheckedContinuation call
let unsafe = self._state.unsafe
unsafe.lock()
let action = unsafe.withValueAssumingLockIsAcquired {
let action = self._state.withLockedValue {
$0.stateMachine.yield(yieldID: yieldID)
}
switch action {
case .callDidYield(let delegate):
// We are allocating a new Deque for every write here
unsafe.unlock()
delegate.didYield(element)
self.unbufferQueuedEvents()
return .yielded
case .throwError(let error):
unsafe.unlock()
throw error
case .suspendTask:
// 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
//
// Dropping and reacquiring the lock may result in yields being reordered but
// only from the perspective of when this function was entered. For example:
//
// - T1 calls _yield
// - T2 calls _yield
// - T2 returns from _yield
// - T1 returns from _yield
//
// This is fine: the async writer doesn't offer any ordering guarantees for
// calls made from different threads.
//
// Within a thread there is no possibility of re-ordering as the call only
// returns once the write has been yielded.
return try await withCheckedThrowingContinuation {
(continuation: CheckedContinuation<StateMachine.YieldResult, Error>) in
let didSuspend = unsafe.withValueAssumingLockIsAcquired {
$0.stateMachine.yield(continuation: continuation, yieldID: yieldID)
return $0.didSuspend
let (action, didSuspend) = self._state.withLockedValue {
state -> (NIOAsyncWriter.StateMachine.YieldAction, (@Sendable () -> Void)?) in
let yieldAction = state.stateMachine.yield(yieldID: yieldID)
switch yieldAction {
case .callDidYield, .throwError:
return (yieldAction, nil)
case .suspendTask:
state.stateMachine.yield(continuation: continuation, yieldID: yieldID)
let didSuspend = state.didSuspend
return (yieldAction, didSuspend)
}
}
switch action {
case .callDidYield(let delegate):
delegate.didYield(element)
self.unbufferQueuedEvents()
continuation.resume(returning: .yielded)
case .throwError(let error):
continuation.resume(throwing: error)
case .suspendTask:
didSuspend?()
}
unsafe.unlock()
didSuspend?()
}
}
} onCancel: {
docs/workarounds.md
+1
View File
@@ -16,3 +16,4 @@ or cases where the Swift compiler was unable to sufficiently optimize the code:
- https://github.com/apple/swift-nio/pull/1961
- https://github.com/apple/swift-nio/pull/2046
- https://github.com/apple/swift-nio/pull/3303
- https://github.com/apple/swift-nio/pull/3452