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

Code before:
/// 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)
  defer { unsafe Builtin.taskRemoveCancellationHandler(record: record) }

Review comment (bug):
do we have any tests that exercise the case [described here](https://github.com/swiftlang/swift/pull/79817#discussion_r2038011762)?

```swift
func next(isolation: isolated (any Actor)? = #isolation) async {
    MainActor.assertIsolated() // passes
    await withTaskCancellationHandler {
        MainActor.assertIsolated() // fails!
    } onCancel: { }
}

@MainActor
func callOnMain() async {
  await next()
}
```

is the outcome in that scenario expected to be different now? or will the same issue occur depending on how the closure isolation is inferred?

Code after:
/// 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)
  defer { unsafe Builtin.taskRemoveCancellationHandler(record: record) }