Deprecation Housekeeping (#206)

`withTaskCancellationHandler<T>(_:, operation:)` -> `withTaskCancellationHandler<T>(operation:onCancel:)`

Author: rustle

Sources/AsyncAlgorithms/AsyncBufferSequence.swift

@@ -265,15 +265,15 @@ extension AsyncBufferSequence: AsyncSequence {
 
       func next() async rethrows -> Element? {
         let result: Result<Element?, Error> = await withTaskCancellationHandler {
-          task?.cancel()
-        } operation: {
           do {
             let value = try await state.next(buffer: buffer)
             return .success(value)
           } catch {
             task?.cancel()
             return .failure(error)
           }
+        } onCancel: {
+          task?.cancel()
         }
         return try result._rethrowGet()
       }

Sources/AsyncAlgorithms/AsyncChannel.swift

@@ -40,10 +40,10 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
       let generation = channel.establish()
       let nextTokenStatus = ManagedCriticalState<ChannelTokenStatus>(.new)
 
-      let value = await withTaskCancellationHandler { [channel] in
-        channel.cancelNext(nextTokenStatus, generation)
-      } operation: {
+      let value = await withTaskCancellationHandler {
         await channel.next(nextTokenStatus, generation)
+      } onCancel: { [channel] in
+        channel.cancelNext(nextTokenStatus, generation)
       }
 
       if let value {
@@ -237,10 +237,10 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
     let generation = establish()
     let sendTokenStatus = ManagedCriticalState<ChannelTokenStatus>(.new)
 
-    await withTaskCancellationHandler { [weak self] in
-      self?.cancelSend(sendTokenStatus, generation)
-    } operation: {
+    await withTaskCancellationHandler {
       await send(sendTokenStatus, generation, element)
+    } onCancel: { [weak self] in
+      self?.cancelSend(sendTokenStatus, generation)
     }
   }
 

Sources/AsyncAlgorithms/AsyncThrowingChannel.swift

@@ -39,10 +39,10 @@ public final class AsyncThrowingChannel<Element: Sendable, Failure: Error>: Asyn
       let nextTokenStatus = ManagedCriticalState<ChannelTokenStatus>(.new)
 
       do {
-        let value = try await withTaskCancellationHandler { [channel] in
-          channel.cancelNext(nextTokenStatus, generation)
-        } operation: {
+        let value = try await withTaskCancellationHandler {
           try await channel.next(nextTokenStatus, generation)
+        } onCancel: { [channel] in
+          channel.cancelNext(nextTokenStatus, generation)
         }
 
         if let value = value {
@@ -295,10 +295,10 @@ public final class AsyncThrowingChannel<Element: Sendable, Failure: Error>: Asyn
     let generation = establish()
     let sendTokenStatus = ManagedCriticalState<ChannelTokenStatus>(.new)
 
-    await withTaskCancellationHandler { [weak self] in
-      self?.cancelSend(sendTokenStatus, generation)
-    } operation: {
+    await withTaskCancellationHandler {
       await send(sendTokenStatus, generation, element)
+    } onCancel: { [weak self] in
+      self?.cancelSend(sendTokenStatus, generation)
     }
   }
 

Sources/AsyncAlgorithms/TaskSelect.swift

@@ -36,14 +36,6 @@ extension Task {
   where Tasks.Element == Task<Success, Failure> {
     let state = ManagedCriticalState(TaskSelectState<Success, Failure>())
     return await withTaskCancellationHandler {
-      let tasks = state.withCriticalRegion { state -> [Task<Success, Failure>] in
-        defer { state.tasks = nil }
-        return state.tasks ?? []
-      }
-      for task in tasks {
-        task.cancel()
-      }
-    } operation: {
       await withUnsafeContinuation { continuation in
         for task in tasks {
           Task<Void, Never> {
@@ -61,6 +53,14 @@ extension Task {
           }?.cancel()
         }
       }
+    } onCancel: {
+      let tasks = state.withCriticalRegion { state -> [Task<Success, Failure>] in
+        defer { state.tasks = nil }
+        return state.tasks ?? []
+      }
+      for task in tasks {
+        task.cancel()
+      }
     }
   }
 

Sources/AsyncSequenceValidation/Clock.swift

@@ -114,11 +114,11 @@ extension AsyncSequenceValidationDiagram.Clock {
   ) async throws {
     let token = queue.prepare()
     try await withTaskCancellationHandler {
-      queue.cancel(token)
-    } operation: {
       try await withUnsafeThrowingContinuation { continuation in
         queue.enqueue(AsyncSequenceValidationDiagram.Context.currentJob, deadline: deadline, continuation: continuation, token: token)
       }
+    } onCancel: {
+      queue.cancel(token)
     }
   }
 }

Sources/AsyncSequenceValidation/Input.swift

@@ -50,12 +50,12 @@ extension AsyncSequenceValidationDiagram {
             eventIndex = 0
           }
         }
-        return try await withTaskCancellationHandler { [queue] in
-          queue.cancel(token)
-        } operation: {
+        return try await withTaskCancellationHandler {
           try await withUnsafeThrowingContinuation { continuation in
             queue.enqueue(Context.currentJob, deadline: when, continuation: continuation, results[eventIndex], index: index, token: token)
           }
+        } onCancel: { [queue] in
+          queue.cancel(token)
         }
       }
 

Tests/AsyncAlgorithmsTests/Support/ManualClock.swift

@@ -245,11 +245,11 @@ public struct ManualClock: Clock {
       return state.generation
     }
     try await withTaskCancellationHandler {
-      cancel(generation)
-    } operation: {
       try await withUnsafeThrowingContinuation { continuation in
         schedule(generation, continuation: continuation, deadline: deadline)
       }
+    } onCancel: {
+      cancel(generation)
     }
   }
 }

Tests/AsyncAlgorithmsTests/TestTaskSelect.swift

@@ -67,28 +67,28 @@ final class TestTaskSelect: XCTestCase {
     let secondCancelled = expectation(description: "second cancelled")
     let task = Task {
       _ = await Task.select(Task {
-        await withTaskCancellationHandler {
-          firstCancelled.fulfill()
-        } operation: { () -> Int in
+        await withTaskCancellationHandler { () -> Int in
           firstReady.fulfill()
           if #available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *) {
             try? await Task.sleep(until: .now + .seconds(2), clock: .continuous)
           } else {
             try? await Task.sleep(nanoseconds: 2_000_000_000)
           }
           return 1
+        } onCancel: {
+          firstCancelled.fulfill()
         }
       }, Task {
-        await withTaskCancellationHandler {
-          secondCancelled.fulfill()
-        } operation: { () -> Int in
+        await withTaskCancellationHandler { () -> Int in
           secondReady.fulfill()
           if #available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *) {
             try? await Task.sleep(until: .now + .seconds(2), clock: .continuous)
           } else {
             try? await Task.sleep(nanoseconds: 2_000_000_000)
           }
           return 1
+        } onCancel: {
+          secondCancelled.fulfill()
         }
       })
     }