[Concurrency] shim for deprecated Task.withCancellationHandler

Author: ktoso

stdlib/public/Concurrency/TaskCancellation.swift

@@ -15,6 +15,37 @@ import Swift
 
 // ==== Task Cancellation ------------------------------------------------------
 
+/// Execute an operation with cancellation handler which will immediately be
+/// invoked if the current task is cancelled.
+///
+/// This differs from the operation cooperatively checking for cancellation
+/// and reacting to it in that the cancellation handler is _always_ and
+/// _immediately_ invoked when the task is cancelled. For example, even if the
+/// operation is running code which never checks for cancellation, a cancellation
+/// handler still would run and give us a chance to run some cleanup code.
+///
+/// Does not check for cancellation, and always executes the passed `operation`.
+///
+/// This function returns instantly and will never suspend.
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+public func withTaskCancellationHandler<T>(
+  handler: @Sendable () -> (),
+  operation: () async throws -> T
+) async rethrows -> T {
+  let task = Builtin.getCurrentAsyncTask()
+
+  guard !_taskIsCancelled(task) else {
+    // If the current task is already cancelled, run the handler immediately.
+    handler()
+    return try await operation()
+  }
+
+  let record = _taskAddCancellationHandler(handler: handler)
+  defer { _taskRemoveCancellationHandler(record: record) }
+
+  return try await operation()
+}
+
 @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
 extension Task {
 
@@ -57,34 +88,12 @@ extension Task {
     }
   }
 
-  /// Execute an operation with cancellation handler which will immediately be
-  /// invoked if the current task is cancelled.
-  ///
-  /// This differs from the operation cooperatively checking for cancellation
-  /// and reacting to it in that the cancellation handler is _always_ and
-  /// _immediately_ invoked when the task is cancelled. For example, even if the
-  /// operation is running code which never checks for cancellation, a cancellation
-  /// handler still would run and give us a chance to run some cleanup code.
-  ///
-  /// Does not check for cancellation, and always executes the passed `operation`.
-  ///
-  /// This function returns instantly and will never suspend.
+  @available(*, deprecated, message: "`Task.withCancellationHandler` has been replaced by `withTaskCancellationHandler` and will be removed shortly.")
   public static func withCancellationHandler<T>(
     handler: @Sendable () -> (),
     operation: () async throws -> T
   ) async rethrows -> T {
-    let task = Builtin.getCurrentAsyncTask()
-
-    guard !_taskIsCancelled(task) else {
-      // If the current task is already cancelled, run the handler immediately.
-      handler()
-      return try await operation()
-    }
-
-    let record = _taskAddCancellationHandler(handler: handler)
-    defer { _taskRemoveCancellationHandler(record: record) }
-
-    return try await operation()
+    try await withTaskCancellationHandler(handler: handler, operation: operation)
   }
 
   /// The default cancellation thrown when a task is cancelled.
@@ -95,7 +104,6 @@ extension Task {
     // no extra information, cancellation is intended to be light-weight
     public init() {}
   }
-
 }
 
 @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)

test/Concurrency/async_cancellation.swift

@@ -26,15 +26,14 @@ struct SomeFile: Sendable {
 }
 
 @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
-func test_cancellation_withCancellationHandler(_ anything: Any) async -> PictureData {
+func test_cancellation_withTaskCancellationHandler(_ anything: Any) async -> PictureData {
   let handle: Task.Handle<PictureData, Error> = detach {
     let file = SomeFile()
 
-    return await Task.withCancellationHandler(
-      handler: { file.close() },
-      operation: {
+    return await withTaskCancellationHandler(
+      handler: { file.close() }) {
       await test_cancellation_guard_isCancelled(file)
-    })
+    }
   }
 
   handle.cancel()