Skip to content

Commit 39d60ac

Browse files
mbrandonwmluisbrown
mbrandonw
authored and
mluisbrown
committed
Fix task cancellation leak (#1418)
* Fix leak in withTaskCancellation(id:) * Add failing test that is fixed * wip * clean up * update tests * wip * wip * wip * wip * clean up * fix 13.4 * try to fix test (cherry picked from commit bc6f87dc44d85cccc74e49b017557577ecc44ca4) # Conflicts: # Sources/ComposableArchitecture/Effects/Cancellation.swift # Tests/ComposableArchitectureTests/EffectRunTests.swift # Tests/ComposableArchitectureTests/EffectTaskTests.swift
1 parent b7cc644 commit 39d60ac

File tree

4 files changed

+223
-210
lines changed

4 files changed

+223
-210
lines changed

‎Sources/ComposableArchitecture/Effects/Cancellation.swift

Lines changed: 16 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -208,27 +208,24 @@ public func withTaskCancellation<T: Sendable>(
208208
cancelInFlight: Bool = false,
209209
operation: @Sendable @escaping () async throws -> T
210210
) async rethrows -> T {
211-
let task = { () -> Task<T, Error> in
212-
cancellablesLock.lock()
213-
let id = CancelToken(id: id)
211+
let id = CancelToken(id: id)
212+
let (cancellable, task) = cancellablesLock.sync { () -> (AnyDisposable, Task<T, Error>) in
214213
if cancelInFlight {
215214
cancellationCancellables[id]?.forEach { $0.dispose() }
216215
}
217216
let task = Task { try await operation() }
218-
var cancellable: AnyDisposable!
219-
cancellable = AnyDisposable {
220-
task.cancel()
221-
cancellablesLock.sync {
222-
cancellationCancellables[id]?.remove(cancellable)
223-
if cancellationCancellables[id]?.isEmpty == .some(true) {
224-
cancellationCancellables[id] = nil
225-
}
217+
let cancellable = AnyDisposable { task.cancel() }
218+
cancellationCancellables[id, default: []].insert(cancellable)
219+
return (cancellable, task)
220+
}
221+
defer {
222+
cancellablesLock.sync {
223+
cancellationCancellables[id]?.remove(cancellable)
224+
if cancellationCancellables[id]?.isEmpty == .some(true) {
225+
cancellationCancellables[id] = nil
226226
}
227227
}
228-
cancellationCancellables[id, default: []].insert(cancellable)
229-
cancellablesLock.unlock()
230-
return task
231-
}()
228+
}
232229
do {
233230
return try await task.cancellableValue
234231
} catch {
@@ -264,10 +261,8 @@ extension Task where Success == Never, Failure == Never {
264261
/// Cancel any currently in-flight operation with the given identifier.
265262
///
266263
/// - Parameter id: An identifier.
267-
public static func cancel<ID: Hashable & Sendable>(id: ID) async {
268-
await MainActor.run {
269-
cancellablesLock.sync { cancellationCancellables[.init(id: id)]?.forEach { $0.dispose() } }
270-
}
264+
public static func cancel<ID: Hashable & Sendable>(id: ID) {
265+
cancellablesLock.sync { cancellationCancellables[.init(id: id)]?.forEach { $0.dispose() } }
271266
}
272267

273268
/// Cancel any currently in-flight operation with the given identifier.
@@ -276,8 +271,8 @@ extension Task where Success == Never, Failure == Never {
276271
/// identifier.
277272
///
278273
/// - Parameter id: A unique type identifying the operation.
279-
public static func cancel(id: Any.Type) async {
280-
await self.cancel(id: ObjectIdentifier(id))
274+
public static func cancel(id: Any.Type) {
275+
self.cancel(id: ObjectIdentifier(id))
281276
}
282277
}
283278

‎Tests/ComposableArchitectureTests/EffectRunTests.swift

Lines changed: 93 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -8,120 +8,120 @@ import XCTest
88

99
// `@MainActor` introduces issues gathering tests on Linux
1010
#if !os(Linux)
11-
@MainActor
12-
final class EffectRunTests: XCTestCase {
13-
func testRun() async {
14-
struct State: Equatable {}
15-
enum Action: Equatable { case tapped, response }
16-
let reducer = Reducer<State, Action, Void> { state, action, _ in
17-
switch action {
18-
case .tapped:
19-
return .run { send in await send(.response) }
20-
case .response:
21-
return .none
22-
}
11+
@MainActor
12+
final class EffectRunTests: XCTestCase {
13+
func testRun() async {
14+
struct State: Equatable {}
15+
enum Action: Equatable { case tapped, response }
16+
let reducer = Reducer<State, Action, Void> { state, action, _ in
17+
switch action {
18+
case .tapped:
19+
return .run { send in await send(.response) }
20+
case .response:
21+
return .none
2322
}
24-
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
25-
await store.send(.tapped)
26-
await store.receive(.response)
2723
}
24+
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
25+
await store.send(.tapped)
26+
await store.receive(.response)
27+
}
2828

29-
func testRunCatch() async {
30-
struct State: Equatable {}
31-
enum Action: Equatable { case tapped, response }
32-
let reducer = Reducer<State, Action, Void> { state, action, _ in
33-
switch action {
34-
case .tapped:
35-
return .run { _ in
36-
struct Failure: Error {}
37-
throw Failure()
38-
} catch: { @Sendable _, send in // NB: Explicit '@Sendable' required in 5.5.2
39-
await send(.response)
40-
}
41-
case .response:
42-
return .none
29+
func testRunCatch() async {
30+
struct State: Equatable {}
31+
enum Action: Equatable { case tapped, response }
32+
let reducer = Reducer<State, Action, Void> { state, action, _ in
33+
switch action {
34+
case .tapped:
35+
return .run { _ in
36+
struct Failure: Error {}
37+
throw Failure()
38+
} catch: { @Sendable _, send in // NB: Explicit '@Sendable' required in 5.5.2
39+
await send(.response)
4340
}
41+
case .response:
42+
return .none
4443
}
45-
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
46-
await store.send(.tapped)
47-
await store.receive(.response)
4844
}
45+
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
46+
await store.send(.tapped)
47+
await store.receive(.response)
48+
}
4949

5050
// `XCTExpectFailure` is not supported on Linux
5151
#if !os(Linux)
52-
func testRunUnhandledFailure() async {
53-
XCTExpectFailure(nil, enabled: nil, strict: nil) {
54-
$0.compactDescription == """
55-
An 'Effect.run' returned from "ComposableArchitectureTests/EffectRunTests.swift:69" threw \
56-
an unhandled error. …
52+
func testRunUnhandledFailure() async {
53+
XCTExpectFailure(nil, enabled: nil, strict: nil) {
54+
$0.compactDescription == """
55+
An 'Effect.run' returned from "ComposableArchitectureTests/EffectRunTests.swift:69" threw \
56+
an unhandled error. …
5757
58-
EffectRunTests.Failure()
58+
EffectRunTests.Failure()
5959
60-
All non-cancellation errors must be explicitly handled via the 'catch' parameter on \
61-
'Effect.run', or via a 'do' block.
62-
"""
63-
}
64-
struct State: Equatable {}
65-
enum Action: Equatable { case tapped, response }
66-
let reducer = Reducer<State, Action, Void> { state, action, _ in
67-
switch action {
68-
case .tapped:
69-
return .run { send in
70-
struct Failure: Error {}
71-
throw Failure()
72-
}
73-
case .response:
74-
return .none
75-
}
60+
All non-cancellation errors must be explicitly handled via the 'catch' parameter on \
61+
'Effect.run', or via a 'do' block.
62+
"""
63+
}
64+
struct State: Equatable {}
65+
enum Action: Equatable { case tapped, response }
66+
let reducer = Reducer<State, Action, Void> { state, action, _ in
67+
switch action {
68+
case .tapped:
69+
return .run { send in
70+
struct Failure: Error {}
71+
throw Failure()
7672
}
77-
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
78-
// NB: We wait a long time here because XCTest failures take a long time to generate
79-
await store.send(.tapped).finish(timeout: 5 * NSEC_PER_SEC)
73+
case .response:
74+
return .none
8075
}
76+
}
77+
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
78+
// NB: We wait a long time here because XCTest failures take a long time to generate
79+
await store.send(.tapped).finish(timeout: 5 * NSEC_PER_SEC)
80+
}
8181
#endif
8282

83-
func testRunCancellation() async {
84-
enum CancelID {}
85-
struct State: Equatable {}
86-
enum Action: Equatable { case tapped, response }
87-
let reducer = Reducer<State, Action, Void> { state, action, _ in
88-
switch action {
89-
case .tapped:
90-
return .run { send in
91-
await Task.cancel(id: CancelID.self)
92-
try Task.checkCancellation()
93-
await send(.response)
94-
}
95-
.cancellable(id: CancelID.self)
96-
case .response:
97-
return .none
83+
func testRunCancellation() async {
84+
enum CancelID {}
85+
struct State: Equatable {}
86+
enum Action: Equatable { case tapped, response }
87+
let reducer = Reducer<State, Action, Void> { state, action, _ in
88+
switch action {
89+
case .tapped:
90+
return .run { send in
91+
Task.cancel(id: CancelID.self)
92+
try Task.checkCancellation()
93+
await send(.response)
9894
}
95+
.cancellable(id: CancelID.self)
96+
case .response:
97+
return .none
9998
}
100-
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
101-
await store.send(.tapped).finish()
10299
}
100+
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
101+
await store.send(.tapped).finish()
102+
}
103103

104-
func testRunCancellationCatch() async {
105-
enum CancelID {}
106-
struct State: Equatable {}
107-
enum Action: Equatable { case tapped, responseA, responseB }
108-
let reducer = Reducer<State, Action, Void> { state, action, _ in
109-
switch action {
110-
case .tapped:
111-
return .run { send in
112-
await Task.cancel(id: CancelID.self)
113-
try Task.checkCancellation()
114-
await send(.responseA)
115-
} catch: { @Sendable _, send in // NB: Explicit '@Sendable' required in 5.5.2
116-
await send(.responseB)
117-
}
118-
.cancellable(id: CancelID.self)
119-
case .responseA, .responseB:
120-
return .none
104+
func testRunCancellationCatch() async {
105+
enum CancelID {}
106+
struct State: Equatable {}
107+
enum Action: Equatable { case tapped, responseA, responseB }
108+
let reducer = Reducer<State, Action, Void> { state, action, _ in
109+
switch action {
110+
case .tapped:
111+
return .run { send in
112+
Task.cancel(id: CancelID.self)
113+
try Task.checkCancellation()
114+
await send(.responseA)
115+
} catch: { @Sendable _, send in // NB: Explicit '@Sendable' required in 5.5.2
116+
await send(.responseB)
121117
}
118+
.cancellable(id: CancelID.self)
119+
case .responseA, .responseB:
120+
return .none
122121
}
123-
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
124-
await store.send(.tapped).finish()
125122
}
123+
let store = TestStore(initialState: State(), reducer: reducer, environment: ())
124+
await store.send(.tapped).finish()
126125
}
126+
}
127127
#endif

0 commit comments

Comments
 (0)