[Concurrency] Use caller execution for withTaskCannelation handler

The problem is that even with the #isolation parameter the non-Sendable
async closure operation _still_ would potentially hop off the caller
isolation. We introduced this change because the plan was the
non-Sendable closure would run on the isolated parameter's isolation,
but that's not actually the case:

Instead, we can use the @execution(caller) on the function and closure,
in order to guarantee there is no hop between those at all, and
developers can trust that adding this cancellation handler will not
cause any unexpected isolation changes and hops. The API was always
documented to not hop as we execute the operation, so this brings the
correct and expected behavior.

resolves rdar://140110775

move to nonisolated(nonsending)

[Concurrency] latest cancellation handler is emit into client; no ABI entry

Author: ktoso

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,71 @@
+// 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
+
+//public func test<T>(_ block: nonisolated(nonsending) () async -> T ) async  {
+//  await block()
+//}
+
+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
+  }
+}
+
+func testMainActorIsolated(isolation: isolated (any Actor)? = #isolation) async {
+  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:_:)