Repository: apple/swift
File: stdlib/public/Concurrency/TaskCancellation.swift (Swift)

Code before:
///
/// This allows the cancellation handler to set some external "canceled" flag
/// that the operation may be *atomically* checking for in order to avoid
/// performing any actual work once the operation gets to run.
///
/// The `operation` closure executes on the calling execution context, and doesn't
/// suspend or change execution context unless code contained within the closure
/// does so. In other words, the potential suspension point of the
/// `withTaskCancellationHandler(operation:onCancel:)` never suspends by itself before
/// executing the operation.
///
/// If cancellation occurs while the operation is running, the cancellation
/// handler executes *concurrently* with the operation.
///
/// ### Cancellation handlers and locks
///
/// Cancellation handlers which acquire locks must take care to avoid deadlock.
/// The cancellation handler may be invoked while holding internal locks
/// associated with the task or other tasks.  Other operations on the task, such
/// as resuming a continuation, may acquire these same internal locks.
/// Therefore, if a cancellation handler must acquire a lock, other code should
/// not cancel tasks or resume continuations while holding that lock.
@available(SwiftStdlib 5.1, *)
@export(implementation)
nonisolated(nonsending)
public func withTaskCancellationHandler<T, E>(
  operation: () async throws(E) -> T,
  onCancel handler: @Sendable () -> Void
) async throws(E) -> T {
  // unconditionally add the cancellation record to the task.
  // if the task was already cancelled, it will be executed right away.
  let record = unsafe Builtin.taskAddCancellationHandler(handler: handler)
  defer { unsafe Builtin.taskRemoveCancellationHandler(record: record) }

  return try await operation()
}

#if !$Embedded
@available(SwiftStdlib 6.0, *)
@usableFromInline
@abi(func withTaskCancellationHandler<T>(
       operation: () async throws -> T,
       onCancel handler: @Sendable () -> Void,
       isolation: isolated (any Actor)?,
     ) async rethrows -> T)
func __abi_withTaskCancellationHandler<T>(
  operation: () async throws -> T,
  onCancel handler: @Sendable () -> Void,
  isolation: isolated (any Actor)? = #isolation
) async rethrows -> T {
  // unconditionally add the cancellation record to the task.

Review comment (bug):
NIT: Could we give these generic types a better name so that they show up nicely in docs or would that be API breaking?
```suggestion
public func withTaskCancellationHandler<Return, Failure: Error>(
```

Code after:
///
/// This allows the cancellation handler to set some external "canceled" flag
/// that the operation may be *atomically* checking for in order to avoid
/// performing any actual work once the operation gets to run.
///
/// The `operation` closure executes on the calling execution context, and doesn't
/// suspend or change execution context unless code contained within the closure
/// does so. In other words, the potential suspension point of the
/// `withTaskCancellationHandler(operation:onCancel:)` never suspends by itself before
/// executing the operation.
///
/// If cancellation occurs while the operation is running, the cancellation
/// handler executes *concurrently* with the operation.
///
/// ### Cancellation handlers and locks
///
/// Cancellation handlers which acquire locks must take care to avoid deadlock.
/// The cancellation handler may be invoked while holding internal locks
/// associated with the task or other tasks.  Other operations on the task, such
/// as resuming a continuation, may acquire these same internal locks.
/// Therefore, if a cancellation handler must acquire a lock, other code should
/// not cancel tasks or resume continuations while holding that lock.
@available(SwiftStdlib 5.1, *)
@_alwaysEmitIntoClient
nonisolated(nonsending)
public func withTaskCancellationHandler<Return, Failure>(
  operation: nonisolated(nonsending) () async throws(Failure) -> Return,
  onCancel handler: sending () -> Void
) async throws(Failure) -> Return {
  // unconditionally add the cancellation record to the task.
  // if the task was already cancelled, it will be executed right away.
  let record = unsafe Builtin.taskAddCancellationHandler(handler: handler)
  defer { unsafe Builtin.taskRemoveCancellationHandler(record: record) }
  return try await operation()
}

#if !$Embedded
@available(SwiftStdlib 6.0, *)
@usableFromInline
@abi(func withTaskCancellationHandler<T>(
       operation: () async throws -> T,
       onCancel handler: @Sendable () -> Void,
       isolation: isolated (any Actor)?,
     ) async rethrows -> T)
func __abi_withTaskCancellationHandler<T>(
  operation: () async throws -> T,
  onCancel handler: @Sendable () -> Void,
  isolation: isolated (any Actor)? = #isolation
) async rethrows -> T {
  // unconditionally add the cancellation record to the task.
  // if the task was already cancelled, it will be executed right away.