Merge pull request #85518 from DougGregor/task-cancellation-rethrows-nonisolated-nonsending

[Concurrency] Adopt typed throws and nonisolated(nonsending) in withTaskCancellationHandler

Author: DougGregor

stdlib/public/Concurrency/TaskCancellation.swift

@@ -74,10 +74,28 @@ import Swift
 /// 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
-@backDeployed(before: SwiftStdlib 6.0)
-#endif
-public func withTaskCancellationHandler<T>(
+@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
@@ -93,6 +111,7 @@ public func withTaskCancellationHandler<T>(
 #endif
   return try await operation()
 }
+#endif
 
 // Note: hack to stage out @_unsafeInheritExecutor forms of various functions
 // in favor of #isolation. The _unsafeInheritExecutor_ prefix is meaningful

stdlib/public/Observation/Sources/Observation/Observations.swift

@@ -240,7 +240,7 @@ public struct Observations<Element: Sendable, Failure: Error>: AsyncSequence, Se
           }, onCancel: {
             // ensure to clean out our continuation uon cancellation
             State.cancel(state, id: id)
-          }, isolation: iterationIsolation)
+          })
           return try await trackEmission(isolation: iterationIsolation, state: state, id: id)
         }
       } catch {

test/Concurrency/actor_withCancellationHandler.swift

@@ -38,7 +38,7 @@ actor Container {
       num += 1 // no error, this runs synchronously on caller context
     } onCancel: {
       // this should error because cancellation is invoked concurrently
-      num += 10 // expected-error{{actor-isolated property 'num' can not be mutated from a Sendable closure}}
+      num += 10 // expected-error{{actor-isolated property 'num' can not be mutated from a nonisolated context}}
     }
   }
 }

test/Concurrency/async_cancellation.swift

@@ -28,15 +28,23 @@ struct SomeFile: Sendable {
   func close() {}
 }
 
+enum HomeworkError: Error {
+case dogAteIt
+}
+
 @available(SwiftStdlib 5.1, *)
 func test_cancellation_withTaskCancellationHandler(_ anything: Any) async -> PictureData? {
   let handle: Task<PictureData, Error> = .init {
     let file = SomeFile()
 
-    return await withTaskCancellationHandler {
-      await test_cancellation_guard_isCancelled(file)
-    } onCancel: {
-      file.close()
+    do throws(HomeworkError) {
+      return try await withTaskCancellationHandler { () throws(HomeworkError) in
+        await test_cancellation_guard_isCancelled(file)
+      } onCancel: {
+        file.close()
+      }
+    } catch .dogAteIt {
+      return PictureData.value("...")
     }
   }
 

test/api-digester/stability-concurrency-abi.test

@@ -89,8 +89,8 @@ Func withCheckedThrowingContinuation(function:_:) has parameter 0 type change fr
 Func withCheckedThrowingContinuation(function:_:) has parameter 1 type change from (_Concurrency.CheckedContinuation<τ_0_0, any Swift.Error>) -> () to Swift.String
 
 // #isolation adoption for cancellation handlers; old APIs are kept ABI compatible
-Func withTaskCancellationHandler(operation:onCancel:) has been renamed to Func withTaskCancellationHandler(operation:onCancel:isolation:)
-Func withTaskCancellationHandler(operation:onCancel:) has mangled name changing from '_Concurrency.withTaskCancellationHandler<A>(operation: () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A' to '_Concurrency.withTaskCancellationHandler<A>(operation: () async throws -> A, onCancel: @Sendable () -> (), isolation: isolated Swift.Optional<Swift.Actor>) async throws -> A'
+Func withTaskCancellationHandler(operation:onCancel:) has been renamed to Func __abi_withTaskCancellationHandler(operation:onCancel:isolation:)
+Func withTaskCancellationHandler(operation:onCancel:) has mangled name changing from '_Concurrency.withTaskCancellationHandler<A>(operation: () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A' to '_Concurrency.__abi_withTaskCancellationHandler<A>(operation: () async throws -> A, onCancel: @Sendable () -> (), isolation: isolated Swift.Optional<Swift.Actor>) async throws -> A'
 
 // #isolated was adopted and the old methods kept: $ss31withCheckedThrowingContinuation8function_xSS_yScCyxs5Error_pGXEtYaKlF
 Func withCheckedContinuation(function:_:) has been renamed to Func withCheckedContinuation(isolation:function:_:)