update for new withTaskCancellationHandler parameter order

lhoward · lhoward · commit 07518fe18749 · 2024-11-21T14:19:25.000+11:00
diff --git a/Sources/AsyncChannels/AsyncBufferedChannel.swift b/Sources/AsyncChannels/AsyncBufferedChannel.swift
@@ -157,27 +157,7 @@ public final class AsyncBufferedChannel<Element>: AsyncSequence, Sendable {
     let awaitingId = self.generateId()
     let cancellation = ManagedCriticalState<Bool>(false)
 
-    return await withTaskCancellationHandler { [state] in
-      let awaiting = state.withCriticalRegion { state -> Awaiting? in
-        cancellation.withCriticalRegion { cancellation in
-          cancellation = true
-        }
-        switch state {
-          case .awaiting(var awaitings):
-            let awaiting = awaitings.remove(.placeHolder(id: awaitingId))
-            if awaitings.isEmpty {
-              state = .idle
-            } else {
-              state = .awaiting(awaitings)
-            }
-            return awaiting
-          default:
-            return nil
-        }
-      }
-
-      awaiting?.continuation?.resume(returning: nil)
-    } operation: {
+    return await withTaskCancellationHandler {
       await withCheckedContinuation { [state] (continuation: CheckedContinuation<Element?, Never>) in
         let decision = state.withCriticalRegion { state -> AwaitingDecision in
           let isCancelled = cancellation.withCriticalRegion { $0 }
@@ -218,6 +198,26 @@ public final class AsyncBufferedChannel<Element>: AsyncSequence, Sendable {
             onSuspend?()
         }
       }
+    } onCancel: { [state] in
+      let awaiting = state.withCriticalRegion { state -> Awaiting? in
+        cancellation.withCriticalRegion { cancellation in
+          cancellation = true
+        }
+        switch state {
+          case .awaiting(var awaitings):
+            let awaiting = awaitings.remove(.placeHolder(id: awaitingId))
+            if awaitings.isEmpty {
+              state = .idle
+            } else {
+              state = .awaiting(awaitings)
+            }
+            return awaiting
+          default:
+            return nil
+        }
+      }
+
+      awaiting?.continuation?.resume(returning: nil)
     }
   }
 
diff --git a/Sources/AsyncChannels/AsyncThrowingBufferedChannel.swift b/Sources/AsyncChannels/AsyncThrowingBufferedChannel.swift
@@ -178,27 +178,7 @@ public final class AsyncThrowingBufferedChannel<Element, Failure: Error>: AsyncS
     let awaitingId = self.generateId()
     let cancellation = ManagedCriticalState<Bool>(false)
 
-    return try await withTaskCancellationHandler { [state] in
-      let awaiting = state.withCriticalRegion { state -> Awaiting? in
-        cancellation.withCriticalRegion { cancellation in
-          cancellation = true
-        }
-        switch state {
-          case .awaiting(var awaitings):
-            let awaiting = awaitings.remove(.placeHolder(id: awaitingId))
-            if awaitings.isEmpty {
-              state = .idle
-            } else {
-              state = .awaiting(awaitings)
-            }
-            return awaiting
-          default:
-            return nil
-        }
-      }
-
-      awaiting?.continuation?.resume(returning: nil)
-    } operation: {
+    return try await withTaskCancellationHandler {
       try await withUnsafeThrowingContinuation { [state] (continuation: UnsafeContinuation<Element?, Error>) in
         let decision = state.withCriticalRegion { state -> AwaitingDecision in
           let isCancelled = cancellation.withCriticalRegion { $0 }
@@ -245,6 +225,26 @@ public final class AsyncThrowingBufferedChannel<Element, Failure: Error>: AsyncS
             onSuspend?()
         }
       }
+    } onCancel: { [state] in
+      let awaiting = state.withCriticalRegion { state -> Awaiting? in
+        cancellation.withCriticalRegion { cancellation in
+          cancellation = true
+        }
+        switch state {
+          case .awaiting(var awaitings):
+            let awaiting = awaitings.remove(.placeHolder(id: awaitingId))
+            if awaitings.isEmpty {
+              state = .idle
+            } else {
+              state = .awaiting(awaitings)
+            }
+            return awaiting
+          default:
+            return nil
+        }
+      }
+
+      awaiting?.continuation?.resume(returning: nil)
     }
   }
 
diff --git a/Sources/AsyncSubjects/AsyncCurrentValueSubject.swift b/Sources/AsyncSubjects/AsyncCurrentValueSubject.swift
@@ -137,10 +137,10 @@ public final class AsyncCurrentValueSubject<Element>: AsyncSubject where Element
     }
 
     public mutating func next() async -> Element? {
-      await withTaskCancellationHandler { [unregister] in
-        unregister()
-      } operation: {
+      await withTaskCancellationHandler {
         await self.iterator.next()
+      } onCancel: { [unregister] in
+        unregister()
       }
     }
   }
diff --git a/Sources/AsyncSubjects/AsyncPassthroughSubject.swift b/Sources/AsyncSubjects/AsyncPassthroughSubject.swift
@@ -120,10 +120,10 @@ public final class AsyncPassthroughSubject<Element>: AsyncSubject {
     }
 
     public mutating func next() async -> Element? {
-      await withTaskCancellationHandler { [unregister] in
-        unregister()
-      } operation: {
+      await withTaskCancellationHandler {
         await self.iterator.next()
+      } onCancel: { [unregister] in
+        unregister()
       }
     }
   }
diff --git a/Sources/AsyncSubjects/AsyncReplaySubject.swift b/Sources/AsyncSubjects/AsyncReplaySubject.swift
@@ -124,10 +124,10 @@ public final class AsyncReplaySubject<Element>: AsyncSubject where Element: Send
     }
 
     public mutating func next() async -> Element? {
-      await withTaskCancellationHandler { [unregister] in
-        unregister()
-      } operation: {
+      await withTaskCancellationHandler {
         await self.iterator.next()
+      } onCancel: { [unregister] in
+        unregister()
       }
     }
   }
diff --git a/Sources/AsyncSubjects/AsyncThrowingCurrentValueSubject.swift b/Sources/AsyncSubjects/AsyncThrowingCurrentValueSubject.swift
@@ -148,10 +148,10 @@ public final class AsyncThrowingCurrentValueSubject<Element, Failure: Error>: As
     }
 
     public mutating func next() async throws -> Element? {
-      try await withTaskCancellationHandler { [unregister] in
-        unregister()
-      } operation: {
+      try await withTaskCancellationHandler {
         try await self.iterator.next()
+      } onCancel: { [unregister] in
+        unregister()
       }
     }
   }
diff --git a/Sources/AsyncSubjects/AsyncThrowingPassthroughSubject.swift b/Sources/AsyncSubjects/AsyncThrowingPassthroughSubject.swift
@@ -132,10 +132,10 @@ public final class AsyncThrowingPassthroughSubject<Element, Failure: Error>: Asy
     }
 
     public mutating func next() async throws -> Element? {
-      try await withTaskCancellationHandler { [unregister] in
-        unregister()
-      } operation: {
+      try await withTaskCancellationHandler {
         try await self.iterator.next()
+      } onCancel: { [unregister] in
+        unregister()
       }
     }
   }
diff --git a/Sources/AsyncSubjects/AsyncThrowingReplaySubject.swift b/Sources/AsyncSubjects/AsyncThrowingReplaySubject.swift
@@ -134,10 +134,10 @@ public final class AsyncThrowingReplaySubject<Element, Failure: Error>: AsyncSub
     }
 
     public mutating func next() async throws -> Element? {
-      try await withTaskCancellationHandler { [unregister] in
-        unregister()
-      } operation: {
+      try await withTaskCancellationHandler {
         try await self.iterator.next()
+      } onCancel: { [unregister] in
+        unregister()
       }
     }
   }
diff --git a/Sources/Combiners/Merge/MergeStateMachine.swift b/Sources/Combiners/Merge/MergeStateMachine.swift
@@ -196,8 +196,6 @@ struct MergeStateMachine<Element>: Sendable {
 
   func next() async -> RegulatedElement<Element> {
     await withTaskCancellationHandler {
-      self.unsuspendAndClearOnCancel()
-    } operation: {
       self.requestNextRegulatedElements()
 
       let regulatedElement = await withUnsafeContinuation { (continuation: UnsafeContinuation<RegulatedElement<Element>, Never>) in
@@ -244,6 +242,8 @@ struct MergeStateMachine<Element>: Sendable {
       }
 
       return regulatedElement
+    } onCancel: {
+      self.unsuspendAndClearOnCancel()
     }
   }
 }
diff --git a/Sources/Combiners/WithLatestFrom/AsyncWithLatestFrom2Sequence.swift b/Sources/Combiners/WithLatestFrom/AsyncWithLatestFrom2Sequence.swift
@@ -172,12 +172,7 @@ where Other1: Sendable, Other2: Sendable, Other1.Element: Sendable, Other2.Eleme
       let shouldReturnNil = self.isTerminated.withCriticalRegion { $0 }
       guard !shouldReturnNil else { return nil }
 
-      return try await withTaskCancellationHandler { [isTerminated, othersTask] in
-        isTerminated.withCriticalRegion { isTerminated in
-          isTerminated = true
-        }
-        othersTask?.cancel()
-      } operation: { [othersTask, othersState, onBaseElement] in
+      return try await withTaskCancellationHandler { [othersTask, othersState, onBaseElement] in
         do {
           while true {
             guard let baseElement = try await self.base.next() else {
@@ -219,6 +214,11 @@ where Other1: Sendable, Other2: Sendable, Other1.Element: Sendable, Other2.Eleme
           othersTask?.cancel()
           throw error
         }
+      } onCancel: { [isTerminated, othersTask] in
+        isTerminated.withCriticalRegion { isTerminated in
+          isTerminated = true
+        }
+        othersTask?.cancel()
       }
     }
   }
diff --git a/Sources/Combiners/WithLatestFrom/AsyncWithLatestFromSequence.swift b/Sources/Combiners/WithLatestFrom/AsyncWithLatestFromSequence.swift
@@ -121,9 +121,7 @@ where Other: Sendable, Other.Element: Sendable {
     public mutating func next() async rethrows -> Element? {
       guard !self.isTerminated else { return nil }
 
-      return try await withTaskCancellationHandler { [otherTask] in
-        otherTask?.cancel()
-      } operation: { [otherTask, otherState, onBaseElement] in
+      return try await withTaskCancellationHandler { [otherTask, otherState, onBaseElement] in
         do {
           while true {
             guard let baseElement = try await self.base.next() else {
@@ -157,8 +155,10 @@ where Other: Sendable, Other.Element: Sendable {
           otherTask?.cancel()
           throw error
         }
+      } onCancel: { [otherTask] in
+        otherTask?.cancel()
       }
-    }
+   }
   }
 }
 
diff --git a/Sources/Creators/AsyncTimerSequence.swift b/Sources/Creators/AsyncTimerSequence.swift
@@ -78,11 +78,11 @@ public struct AsyncTimerSequence: AsyncSequence {
     }
 
     public mutating func next() async -> Element? {
-      await withTaskCancellationHandler { [task] in
-        task.cancel()
-      } operation: {
+      await withTaskCancellationHandler {
         guard !Task.isCancelled else { return nil }
         return await self.iterator.next()
+      } onCancel: { [task] in
+        task.cancel()
       }
     }
   }
diff --git a/Sources/Supporting/Regulator.swift b/Sources/Supporting/Regulator.swift
@@ -48,8 +48,6 @@ final class Regulator<Base: AsyncSequence>: @unchecked Sendable {
 
   func iterate() async {
     await withTaskCancellationHandler {
-      self.unsuspendAndExitOnCancel()
-    } operation: {
       var mutableBase = base.makeAsyncIterator()
 
       do {
@@ -99,6 +97,8 @@ final class Regulator<Base: AsyncSequence>: @unchecked Sendable {
         }
         self.onNextRegulatedElement(.element(result: .failure(error)))
       }
+    } onCancel: {
+      self.unsuspendAndExitOnCancel()
     }
   }
 
diff --git a/Tests/Operators/AsyncSequence+ShareTests.swift b/Tests/Operators/AsyncSequence+ShareTests.swift
@@ -40,15 +40,15 @@ private struct LongAsyncSequence<Element>: AsyncSequence, AsyncIteratorProtocol
   }
   
   mutating func next() async throws -> Element? {
-    return try await withTaskCancellationHandler { [onCancel] in
-      onCancel()
-    } operation: {
+    return try await withTaskCancellationHandler {
       try await Task.sleep(nanoseconds: self.interval.nanoseconds)
       self.currentIndex += 1
       if self.currentIndex == self.failAt {
         throw MockError(code: 0)
       }
       return self.elements.next()
+    } onCancel: { [onCancel] in
+      onCancel()
     }
   }
   
diff --git a/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift b/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift
@@ -41,15 +41,15 @@ private struct LongAsyncSequence<Element>: AsyncSequence, AsyncIteratorProtocol
   }
 
   mutating func next() async throws -> Element? {
-    return try await withTaskCancellationHandler { [onCancel] in
-      onCancel()
-    } operation: {
+    return try await withTaskCancellationHandler {
       try await Task.sleep(nanoseconds: self.interval.nanoseconds)
       self.currentIndex += 1
       if self.currentIndex == self.failAt {
         throw MockError(code: 0)
       }
       return self.elements.next()
+    } onCancel: { [onCancel] in
+      onCancel()
     }
   }
 

Original file line number	Diff line number	Diff line change
`@@ -137,10 +137,10 @@ public final class AsyncCurrentValueSubject<Element>: AsyncSubject where Element`
`137`	`137`	`}`
`138`	`138`
`139`	`139`	`public mutating func next() async -> Element? {`
`140`		`- await withTaskCancellationHandler { [unregister] in`
`141`		`- unregister()`
`142`		`- } operation: {`
	`140`	`+ await withTaskCancellationHandler {`
`143`	`141`	`await self.iterator.next()`
	`142`	`+ } onCancel: { [unregister] in`
	`143`	`+ unregister()`
`144`	`144`	`}`
`145`	`145`	`}`
`146`	`146`	`}`
Original file line number	Diff line number	Diff line change
`@@ -120,10 +120,10 @@ public final class AsyncPassthroughSubject<Element>: AsyncSubject {`
`120`	`120`	`}`
`121`	`121`
`122`	`122`	`public mutating func next() async -> Element? {`
`123`		`- await withTaskCancellationHandler { [unregister] in`
`124`		`- unregister()`
`125`		`- } operation: {`
	`123`	`+ await withTaskCancellationHandler {`
`126`	`124`	`await self.iterator.next()`
	`125`	`+ } onCancel: { [unregister] in`
	`126`	`+ unregister()`
`127`	`127`	`}`
`128`	`128`	`}`
`129`	`129`	`}`
Original file line number	Diff line number	Diff line change
`@@ -124,10 +124,10 @@ public final class AsyncReplaySubject<Element>: AsyncSubject where Element: Send`
`124`	`124`	`}`
`125`	`125`
`126`	`126`	`public mutating func next() async -> Element? {`
`127`		`- await withTaskCancellationHandler { [unregister] in`
`128`		`- unregister()`
`129`		`- } operation: {`
	`127`	`+ await withTaskCancellationHandler {`
`130`	`128`	`await self.iterator.next()`
	`129`	`+ } onCancel: { [unregister] in`
	`130`	`+ unregister()`
`131`	`131`	`}`
`132`	`132`	`}`
`133`	`133`	`}`
Original file line number	Diff line number	Diff line change
`@@ -148,10 +148,10 @@ public final class AsyncThrowingCurrentValueSubject<Element, Failure: Error>: As`
`148`	`148`	`}`
`149`	`149`
`150`	`150`	`public mutating func next() async throws -> Element? {`
`151`		`- try await withTaskCancellationHandler { [unregister] in`
`152`		`- unregister()`
`153`		`- } operation: {`
	`151`	`+ try await withTaskCancellationHandler {`
`154`	`152`	`try await self.iterator.next()`
	`153`	`+ } onCancel: { [unregister] in`
	`154`	`+ unregister()`
`155`	`155`	`}`
`156`	`156`	`}`
`157`	`157`	`}`
Original file line number	Diff line number	Diff line change
`@@ -132,10 +132,10 @@ public final class AsyncThrowingPassthroughSubject<Element, Failure: Error>: Asy`
`132`	`132`	`}`
`133`	`133`
`134`	`134`	`public mutating func next() async throws -> Element? {`
`135`		`- try await withTaskCancellationHandler { [unregister] in`
`136`		`- unregister()`
`137`		`- } operation: {`
	`135`	`+ try await withTaskCancellationHandler {`
`138`	`136`	`try await self.iterator.next()`
	`137`	`+ } onCancel: { [unregister] in`
	`138`	`+ unregister()`
`139`	`139`	`}`
`140`	`140`	`}`
`141`	`141`	`}`
Original file line number	Diff line number	Diff line change
`@@ -134,10 +134,10 @@ public final class AsyncThrowingReplaySubject<Element, Failure: Error>: AsyncSub`
`134`	`134`	`}`
`135`	`135`
`136`	`136`	`public mutating func next() async throws -> Element? {`
`137`		`- try await withTaskCancellationHandler { [unregister] in`
`138`		`- unregister()`
`139`		`- } operation: {`
	`137`	`+ try await withTaskCancellationHandler {`
`140`	`138`	`try await self.iterator.next()`
	`139`	`+ } onCancel: { [unregister] in`
	`140`	`+ unregister()`
`141`	`141`	`}`
`142`	`142`	`}`
`143`	`143`	`}`
Original file line number	Diff line number	Diff line change
`@@ -196,8 +196,6 @@ struct MergeStateMachine<Element>: Sendable {`
`196`	`196`
`197`	`197`	`func next() async -> RegulatedElement<Element> {`
`198`	`198`	`await withTaskCancellationHandler {`
`199`		`- self.unsuspendAndClearOnCancel()`
`200`		`- } operation: {`
`201`	`199`	`self.requestNextRegulatedElements()`
`202`	`200`
`203`	`201`	`let regulatedElement = await withUnsafeContinuation { (continuation: UnsafeContinuation<RegulatedElement<Element>, Never>) in`
`@@ -244,6 +242,8 @@ struct MergeStateMachine<Element>: Sendable {`
`244`	`242`	`}`
`245`	`243`
`246`	`244`	`return regulatedElement`
	`245`	`+ } onCancel: {`
	`246`	`+ self.unsuspendAndClearOnCancel()`
`247`	`247`	`}`
`248`	`248`	`}`
`249`	`249`	`}`
Original file line number	Diff line number	Diff line change
`@@ -172,12 +172,7 @@ where Other1: Sendable, Other2: Sendable, Other1.Element: Sendable, Other2.Eleme`
`172`	`172`	`let shouldReturnNil = self.isTerminated.withCriticalRegion { $0 }`
`173`	`173`	`guard !shouldReturnNil else { return nil }`
`174`	`174`
`175`		`- return try await withTaskCancellationHandler { [isTerminated, othersTask] in`
`176`		`- isTerminated.withCriticalRegion { isTerminated in`
`177`		`- isTerminated = true`
`178`		`- }`
`179`		`- othersTask?.cancel()`
`180`		`- } operation: { [othersTask, othersState, onBaseElement] in`
	`175`	`+ return try await withTaskCancellationHandler { [othersTask, othersState, onBaseElement] in`
`181`	`176`	`do {`
`182`	`177`	`while true {`
`183`	`178`	`guard let baseElement = try await self.base.next() else {`
`@@ -219,6 +214,11 @@ where Other1: Sendable, Other2: Sendable, Other1.Element: Sendable, Other2.Eleme`
`219`	`214`	`othersTask?.cancel()`
`220`	`215`	`throw error`
`221`	`216`	`}`
	`217`	`+ } onCancel: { [isTerminated, othersTask] in`
	`218`	`+ isTerminated.withCriticalRegion { isTerminated in`
	`219`	`+ isTerminated = true`
	`220`	`+ }`
	`221`	`+ othersTask?.cancel()`
`222`	`222`	`}`
`223`	`223`	`}`
`224`	`224`	`}`