Skip to content

Commit da4e36f

Browse files
FranzBuschFranzBusch
authored
Fix potential deadlocks when resuming a continuation while holding a lock (#303)
* Fix potential deadlocks in zip * Fix debounce * Fix combineLatest * Fix Channel * Fix buffer
1 parent 5bbdcc1 commit da4e36f

File tree

6 files changed

+553
-476
lines changed

6 files changed

+553
-476
lines changed

‎Sources/AsyncAlgorithms/Buffer/BoundedBufferStorage.swift

+71-58
Original file line numberDiff line numberDiff line change
@@ -18,34 +18,47 @@ final class BoundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Send
1818

1919
func next() async -> Result<Base.Element, Error>? {
2020
return await withTaskCancellationHandler {
21-
let (shouldSuspend, result) = self.stateMachine.withCriticalRegion { stateMachine -> (Bool, Result<Base.Element, Error>?) in
21+
let action: BoundedBufferStateMachine<Base>.NextAction? = self.stateMachine.withCriticalRegion { stateMachine in
2222
let action = stateMachine.next()
2323
switch action {
2424
case .startTask(let base):
2525
self.startTask(stateMachine: &stateMachine, base: base)
26-
return (true, nil)
26+
return nil
27+
2728
case .suspend:
28-
return (true, nil)
29-
case .returnResult(let producerContinuation, let result):
30-
producerContinuation?.resume()
31-
return (false, result)
29+
return action
30+
case .returnResult:
31+
return action
3232
}
3333
}
3434

35-
if !shouldSuspend {
36-
return result
35+
switch action {
36+
case .startTask:
37+
// We are handling the startTask in the lock already because we want to avoid
38+
// other inputs interleaving while starting the task
39+
fatalError("Internal inconsistency")
40+
41+
case .suspend:
42+
break
43+
44+
case .returnResult(let producerContinuation, let result):
45+
producerContinuation?.resume()
46+
return result
47+
48+
case .none:
49+
break
3750
}
3851

3952
return await withUnsafeContinuation { (continuation: UnsafeContinuation<Result<Base.Element, Error>?, Never>) in
40-
self.stateMachine.withCriticalRegion { stateMachine in
41-
let action = stateMachine.nextSuspended(continuation: continuation)
42-
switch action {
43-
case .none:
44-
break
45-
case .returnResult(let producerContinuation, let result):
46-
producerContinuation?.resume()
47-
continuation.resume(returning: result)
48-
}
53+
let action = self.stateMachine.withCriticalRegion { stateMachine in
54+
stateMachine.nextSuspended(continuation: continuation)
55+
}
56+
switch action {
57+
case .none:
58+
break
59+
case .returnResult(let producerContinuation, let result):
60+
producerContinuation?.resume()
61+
continuation.resume(returning: result)
4962
}
5063
}
5164
} onCancel: {
@@ -68,15 +81,15 @@ final class BoundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Send
6881

6982
if shouldSuspend {
7083
await withUnsafeContinuation { (continuation: UnsafeContinuation<Void, Never>) in
71-
self.stateMachine.withCriticalRegion { stateMachine in
72-
let action = stateMachine.producerSuspended(continuation: continuation)
73-
74-
switch action {
75-
case .none:
76-
break
77-
case .resumeProducer:
78-
continuation.resume()
79-
}
84+
let action = self.stateMachine.withCriticalRegion { stateMachine in
85+
stateMachine.producerSuspended(continuation: continuation)
86+
}
87+
88+
switch action {
89+
case .none:
90+
break
91+
case .resumeProducer:
92+
continuation.resume()
8093
}
8194
}
8295
}
@@ -86,35 +99,35 @@ final class BoundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Send
8699
break loop
87100
}
88101

89-
self.stateMachine.withCriticalRegion { stateMachine in
90-
let action = stateMachine.elementProduced(element: element)
91-
switch action {
92-
case .none:
93-
break
94-
case .resumeConsumer(let continuation, let result):
95-
continuation.resume(returning: result)
96-
}
102+
let action = self.stateMachine.withCriticalRegion { stateMachine in
103+
stateMachine.elementProduced(element: element)
97104
}
98-
}
99-
100-
self.stateMachine.withCriticalRegion { stateMachine in
101-
let action = stateMachine.finish(error: nil)
102105
switch action {
103106
case .none:
104107
break
105-
case .resumeConsumer(let continuation):
106-
continuation?.resume(returning: nil)
108+
case .resumeConsumer(let continuation, let result):
109+
continuation.resume(returning: result)
107110
}
108111
}
112+
113+
let action = self.stateMachine.withCriticalRegion { stateMachine in
114+
stateMachine.finish(error: nil)
115+
}
116+
switch action {
117+
case .none:
118+
break
119+
case .resumeConsumer(let continuation):
120+
continuation?.resume(returning: nil)
121+
}
109122
} catch {
110-
self.stateMachine.withCriticalRegion { stateMachine in
111-
let action = stateMachine.finish(error: error)
112-
switch action {
113-
case .none:
114-
break
115-
case .resumeConsumer(let continuation):
116-
continuation?.resume(returning: .failure(error))
117-
}
123+
let action = self.stateMachine.withCriticalRegion { stateMachine in
124+
stateMachine.finish(error: error)
125+
}
126+
switch action {
127+
case .none:
128+
break
129+
case .resumeConsumer(let continuation):
130+
continuation?.resume(returning: .failure(error))
118131
}
119132
}
120133
}
@@ -123,16 +136,16 @@ final class BoundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Send
123136
}
124137

125138
func interrupted() {
126-
self.stateMachine.withCriticalRegion { stateMachine in
127-
let action = stateMachine.interrupted()
128-
switch action {
129-
case .none:
130-
break
131-
case .resumeProducerAndConsumer(let task, let producerContinuation, let consumerContinuation):
132-
task.cancel()
133-
producerContinuation?.resume()
134-
consumerContinuation?.resume(returning: nil)
135-
}
139+
let action = self.stateMachine.withCriticalRegion { stateMachine in
140+
stateMachine.interrupted()
141+
}
142+
switch action {
143+
case .none:
144+
break
145+
case .resumeProducerAndConsumer(let task, let producerContinuation, let consumerContinuation):
146+
task.cancel()
147+
producerContinuation?.resume()
148+
consumerContinuation?.resume(returning: nil)
136149
}
137150
}
138151

‎Sources/AsyncAlgorithms/Buffer/UnboundedBufferStorage.swift

+55-46
Original file line numberDiff line numberDiff line change
@@ -19,32 +19,41 @@ final class UnboundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Se
1919
func next() async -> Result<Base.Element, Error>? {
2020
return await withTaskCancellationHandler {
2121

22-
let (shouldSuspend, result) = self.stateMachine.withCriticalRegion { stateMachine -> (Bool, Result<Base.Element, Error>?) in
22+
let action: UnboundedBufferStateMachine<Base>.NextAction? = self.stateMachine.withCriticalRegion { stateMachine in
2323
let action = stateMachine.next()
2424
switch action {
2525
case .startTask(let base):
2626
self.startTask(stateMachine: &stateMachine, base: base)
27-
return (true, nil)
27+
return nil
2828
case .suspend:
29-
return (true, nil)
30-
case .returnResult(let result):
31-
return (false, result)
29+
return action
30+
case .returnResult:
31+
return action
3232
}
3333
}
3434

35-
if !shouldSuspend {
36-
return result
35+
switch action {
36+
case .startTask:
37+
// We are handling the startTask in the lock already because we want to avoid
38+
// other inputs interleaving while starting the task
39+
fatalError("Internal inconsistency")
40+
case .suspend:
41+
break
42+
case .returnResult(let result):
43+
return result
44+
case .none:
45+
break
3746
}
3847

3948
return await withUnsafeContinuation { (continuation: UnsafeContinuation<Result<Base.Element, Error>?, Never>) in
40-
self.stateMachine.withCriticalRegion { stateMachine in
41-
let action = stateMachine.nextSuspended(continuation: continuation)
42-
switch action {
43-
case .none:
44-
break
45-
case .resumeConsumer(let result):
46-
continuation.resume(returning: result)
47-
}
49+
let action = self.stateMachine.withCriticalRegion { stateMachine in
50+
stateMachine.nextSuspended(continuation: continuation)
51+
}
52+
switch action {
53+
case .none:
54+
break
55+
case .resumeConsumer(let result):
56+
continuation.resume(returning: result)
4857
}
4958
}
5059
} onCancel: {
@@ -59,35 +68,35 @@ final class UnboundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Se
5968
let task = Task {
6069
do {
6170
for try await element in base {
62-
self.stateMachine.withCriticalRegion { stateMachine in
63-
let action = stateMachine.elementProduced(element: element)
64-
switch action {
65-
case .none:
66-
break
67-
case .resumeConsumer(let continuation, let result):
68-
continuation.resume(returning: result)
69-
}
71+
let action = self.stateMachine.withCriticalRegion { stateMachine in
72+
stateMachine.elementProduced(element: element)
7073
}
71-
}
72-
73-
self.stateMachine.withCriticalRegion { stateMachine in
74-
let action = stateMachine.finish(error: nil)
7574
switch action {
7675
case .none:
7776
break
78-
case .resumeConsumer(let continuation):
79-
continuation?.resume(returning: nil)
77+
case .resumeConsumer(let continuation, let result):
78+
continuation.resume(returning: result)
8079
}
8180
}
81+
82+
let action = self.stateMachine.withCriticalRegion { stateMachine in
83+
stateMachine.finish(error: nil)
84+
}
85+
switch action {
86+
case .none:
87+
break
88+
case .resumeConsumer(let continuation):
89+
continuation?.resume(returning: nil)
90+
}
8291
} catch {
83-
self.stateMachine.withCriticalRegion { stateMachine in
84-
let action = stateMachine.finish(error: error)
85-
switch action {
86-
case .none:
87-
break
88-
case .resumeConsumer(let continuation):
89-
continuation?.resume(returning: .failure(error))
90-
}
92+
let action = self.stateMachine.withCriticalRegion { stateMachine in
93+
stateMachine.finish(error: error)
94+
}
95+
switch action {
96+
case .none:
97+
break
98+
case .resumeConsumer(let continuation):
99+
continuation?.resume(returning: .failure(error))
91100
}
92101
}
93102
}
@@ -96,15 +105,15 @@ final class UnboundedBufferStorage<Base: AsyncSequence>: Sendable where Base: Se
96105
}
97106

98107
func interrupted() {
99-
self.stateMachine.withCriticalRegion { stateMachine in
100-
let action = stateMachine.interrupted()
101-
switch action {
102-
case .none:
103-
break
104-
case .resumeConsumer(let task, let continuation):
105-
task.cancel()
106-
continuation?.resume(returning: nil)
107-
}
108+
let action = self.stateMachine.withCriticalRegion { stateMachine in
109+
stateMachine.interrupted()
110+
}
111+
switch action {
112+
case .none:
113+
break
114+
case .resumeConsumer(let task, let continuation):
115+
task.cancel()
116+
continuation?.resume(returning: nil)
108117
}
109118
}
110119

0 commit comments

Comments
 (0)