[Concurrency] improve cancellation handler to not hop and use caller execution context

stdlib/public/Concurrency/TaskCancellation.swift

@@ -67,10 +67,26 @@ 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, *)
-#if !$Embedded
-@backDeployed(before: SwiftStdlib 6.0)
-#endif
+@_alwaysEmitIntoClient
+nonisolated(nonsending)
 public func withTaskCancellationHandler<T>(
+  operation: nonisolated(nonsending) () async throws -> T,
+  onCancel handler: @Sendable () -> Void
+) 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 _taskAddCancellationHandler(handler: handler)
+  defer { unsafe _taskRemoveCancellationHandler(record: record) }
+
+  return try await operation()
+}
+
+// Note: Deprecated version which would still hop if we did not close over an `isolated` parameter
+// with the operation closure. Instead, we should do what the docs of this method promise - and not hop at all,
+// by using the new nonisolated(nonsending)
+@available(SwiftStdlib 5.1, *)
+@_silgen_name("$ss27withTaskCancellationHandler9operation8onCancel9isolationxxyYaKXE_yyYbXEScA_pSgYitYaKlF")
+public func _isolatedParam_withTaskCancellationHandler<T>(
   operation: () async throws -> T,
   onCancel handler: @Sendable () -> Void,
   isolation: isolated (any Actor)? = #isolation

test/Concurrency/Runtime/cancellation_handler_operation_does_not_hop.swift

@@ -0,0 +1,70 @@
+// RUN: %target-run-simple-swift( -target %target-swift-5.1-abi-triple %import-libdispatch) | %FileCheck %s
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+// UNSUPPORTED: freestanding
+
+actor Canceller {
+  var hello: String = "checking..."
+
+  func testFunc() async {
+    await withTaskCancellationHandler {
+      self.assertIsolated("wat in \(#function)!")
+      print("testFunc.withTaskCancellationHandler") // CHECK: testFunc.withTaskCancellationHandler
+      self.hello = "done!"
+    } onCancel: {
+      // noop
+    }
+
+    // just a simple check to see we executed the closure
+
+    await globalTestFunc()
+  }
+}
+func globalTestFunc(isolation: isolated (any Actor)? = #isolation) async {
+  isolation!.assertIsolated("wat in \(#function)!")
+  await withTaskCancellationHandler {
+     isolation!.assertIsolated("wat in \(#function)!")
+     print("globalTestFunc.withTaskCancellationHandler") // CHECK: globalTestFunc.withTaskCancellationHandler
+  } onCancel: {
+    // noop
+  }
+}
+
+@MainActor
+func testMainActor() async {
+  MainActor.preconditionIsolated("Expected main actor")
+  await withTaskCancellationHandler {
+    MainActor.preconditionIsolated("expected MainActor")
+  } onCancel: {
+    // noop
+  }
+}
+
+// FIXME: rdar://155313349 - nonisolated(nonsending) closure does not pick up isolated parameter when the closure is throwing
+func testMainActorIsolated(isolation: isolated (any Actor)? = #isolation) async {
+  return // FIXME: until rdar://155313349 is fixed
+
+  isolation!.preconditionIsolated("Expected main actor")
+  MainActor.preconditionIsolated("Expected main actor")
+  await withTaskCancellationHandler {
+    print("_unsafeInheritExecutor_withTaskCancellationHandler")
+    MainActor.preconditionIsolated("expected MainActor")
+  } onCancel: {
+    // noop
+  }
+}
+
+_ = await Canceller().testFunc()
+
+_ = await Task { @MainActor in
+  await testMainActor()
+}.value
+
+_ = await Task { @MainActor in
+  await testMainActorIsolated()
+}.value
+
+print("done") // CHECK: done

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

@@ -89,8 +89,9 @@ 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'
+// but ABI checker does not understand the silgen_names on the ABi-compat APIs
+Func withTaskCancellationHandler(operation:onCancel:) has mangled name changing from '_Concurrency.withTaskCancellationHandler<A>(operation: () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A' to '_Concurrency._isolatedParam_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 _isolatedParam_withTaskCancellationHandler(operation:onCancel:isolation:)
 
 // #isolated was adopted and the old methods kept: $ss31withCheckedThrowingContinuation8function_xSS_yScCyxs5Error_pGXEtYaKlF
 Func withCheckedContinuation(function:_:) has been renamed to Func withCheckedContinuation(isolation:function:_:)