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

Code before:
/// 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.
  // if the task was already cancelled, it will be executed right away.
  let record = unsafe Builtin.taskAddCancellationHandler(handler: handler)

Review comment (suggestion):
Can't this closure be `sending` since it is only send to the task/thread that cancels this task so it is never executed concurrently itself just concurrent to the caller of the `withTaskCancellationHandler`.

Code after:
/// 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.
#if $BuiltinConcurrencyStackNesting
  let record = unsafe Builtin.taskAddCancellationHandler(handler: handler)