[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

Author: ktoso

Diff for: stdlib/public/Concurrency/CMakeLists.txt

@@ -63,6 +63,8 @@ endif()
 list(APPEND SWIFT_RUNTIME_CONCURRENCY_SWIFT_FLAGS
   "-enable-experimental-feature"
   "IsolatedAny"
+  "-enable-experimental-feature"
+  "ExecutionAttribute"
   )
 
 list(APPEND SWIFT_RUNTIME_CONCURRENCY_SWIFT_FLAGS "-strict-memory-safety")

Diff for: stdlib/public/Concurrency/TaskCancellation.swift

@@ -68,9 +68,30 @@ import Swift
 /// not cancel tasks or resume continuations while holding that lock.
 @available(SwiftStdlib 5.1, *)
 #if !$Embedded
-@backDeployed(before: SwiftStdlib 6.0)
+@backDeployed(before: SwiftStdlib 6.2)
 #endif
+@execution(caller)
 public func withTaskCancellationHandler<T>(
+  operation: @execution(caller) () 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 @execution(caller)
+@available(SwiftStdlib 5.1, *)
+//#if !$Embedded
+//@backDeployed(before: SwiftStdlib 6.0)
+//#endif
+@_silgen_name("$ss27withTaskCancellationHandler9operation8onCancel9isolationxxyYaKXE_yyYbXEScA_pSgYitYaKlF")
+public func _isolatedParam_withTaskCancellationHandler<T>(
   operation: () async throws -> T,
   onCancel handler: @Sendable () -> Void,
   isolation: isolated (any Actor)? = #isolation

Diff for: test/Concurrency/Runtime/cancellation_handler_operation_does_not_hop.swift

@@ -0,0 +1,60 @@
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking -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 {
+  func testFunc() async {
+    await withTaskCancellationHandler {
+      self.assertIsolated("wat in \(#function)!")
+    } onCancel: {
+      // noop
+    }
+
+    await globalTestFunc()
+  }
+}
+func globalTestFunc(isolation: isolated (any Actor)? = #isolation) async {
+  isolation!.assertIsolated("wat in \(#function)!")
+  await withTaskCancellationHandler {
+     isolation!.assertIsolated("wat in \(#function)!")
+  } 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

Diff for: test/abi/macOS/arm64/concurrency.swift

@@ -320,6 +320,12 @@ Added: _$sScf13checkIsolatedyyFTq
 // withTaskCancellationHandler gains #isolated
 Added: _$ss27withTaskCancellationHandler9operation8onCancel9isolationxxyYaKXE_yyYbXEScA_pSgYitYaKlF
 Added: _$ss27withTaskCancellationHandler9operation8onCancel9isolationxxyYaKXE_yyYbXEScA_pSgYitYaKlFTu
+// withTaskCancellationHandler but with caller execution
+// Swift.withTaskCancellationHandler<A>(operation: @execution(caller) () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A
+// async function pointer to Swift.withTaskCancellationHandler<A>(operation: @execution(caller) () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A
+Added: _$ss27withTaskCancellationHandler9operation8onCancelxxyYaKYCXE_yyYbXEtYaKlF
+Added: _$ss27withTaskCancellationHandler9operation8onCancelxxyYaKYCXE_yyYbXEtYaKlFTu
+
 // TaskGroup.with... APIs gain #isolated
 Added: _$ss23withDiscardingTaskGroup9returning9isolation4bodyxxm_ScA_pSgYixs0bcD0VzYaXEtYalF
 Added: _$ss23withDiscardingTaskGroup9returning9isolation4bodyxxm_ScA_pSgYixs0bcD0VzYaXEtYalFTu

Diff for: test/abi/macOS/x86_64/concurrency.swift

@@ -320,6 +320,12 @@ Added: _$sScf13checkIsolatedyyFTq
 // withTaskCancellationHandler gains #isolated
 Added: _$ss27withTaskCancellationHandler9operation8onCancel9isolationxxyYaKXE_yyYbXEScA_pSgYitYaKlF
 Added: _$ss27withTaskCancellationHandler9operation8onCancel9isolationxxyYaKXE_yyYbXEScA_pSgYitYaKlFTu
+// withTaskCancellationHandler but with caller execution
+// Swift.withTaskCancellationHandler<A>(operation: @execution(caller) () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A
+// async function pointer to Swift.withTaskCancellationHandler<A>(operation: @execution(caller) () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A
+Added: _$ss27withTaskCancellationHandler9operation8onCancelxxyYaKYCXE_yyYbXEtYaKlF
+Added: _$ss27withTaskCancellationHandler9operation8onCancelxxyYaKYCXE_yyYbXEtYaKlFTu
+
 // TaskGroup.with... APIs gain #isolated
 Added: _$ss23withDiscardingTaskGroup9returning9isolation4bodyxxm_ScA_pSgYixs0bcD0VzYaXEtYalF
 Added: _$ss23withDiscardingTaskGroup9returning9isolation4bodyxxm_ScA_pSgYixs0bcD0VzYaXEtYalFTu

Diff for: test/api-digester/stability-concurrency-abi.test

@@ -88,10 +88,6 @@ Func withCheckedThrowingContinuation(function:_:) has mangled name changing from
 Func withCheckedThrowingContinuation(function:_:) has parameter 0 type change from Swift.String to (any _Concurrency.Actor)?
 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'
-
 // #isolated was adopted and the old methods kept: $ss31withCheckedThrowingContinuation8function_xSS_yScCyxs5Error_pGXEtYaKlF
 Func withCheckedContinuation(function:_:) has been renamed to Func withCheckedContinuation(isolation:function:_:)
 Func withCheckedContinuation(function:_:) has mangled name changing from '_Concurrency.withCheckedContinuation<A>(function: Swift.String, _: (Swift.CheckedContinuation<A, Swift.Never>) -> ()) async -> A' to '_Concurrency.withCheckedContinuation<A>(isolation: isolated Swift.Optional<Swift.Actor>, function: Swift.String, _: (Swift.CheckedContinuation<A, Swift.Never>) -> ()) async -> A'
@@ -128,6 +124,11 @@ Func withThrowingTaskGroup(of:returning:body:) has parameter 2 type change from
 Func withTaskGroup(of:returning:body:) has been renamed to Func withTaskGroup(of:returning:isolation:body:)
 Func withTaskGroup(of:returning:body:) has mangled name changing from '_Concurrency.withTaskGroup<A, B where A: Swift.Sendable>(of: A.Type, returning: B.Type, body: (inout Swift.TaskGroup<A>) async -> B) async -> B' to '_Concurrency.withTaskGroup<A, B where A: Swift.Sendable>(of: A.Type, returning: B.Type, isolation: isolated Swift.Optional<Swift.Actor>, body: (inout Swift.TaskGroup<A>) async -> B) async -> B'
 
+// withTaskCancellationHandler now uses caller execution
+// Old methods are kept for ABI compatibility but this test does not understand that
+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: @execution(caller) () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A'
+Func withTaskCancellationHandler(operation:onCancel:) is now with @execution
+
 // *** DO NOT DISABLE OR XFAIL THIS TEST. *** (See comment above.)
 
 

Diff for: utils/swift_snapshot_tool/Package.resolved

@@ -0,0 +1,15 @@
+{
+  "originHash" : "80ba966d30db4d84324c2bda4f2b419b9e6fe208fa1f70080f74c7aa3d432a49",
+  "pins" : [
+    {
+      "identity" : "swift-argument-parser",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-argument-parser",
+      "state" : {
+        "revision" : "41982a3656a71c768319979febd796c6fd111d5c",
+        "version" : "1.5.0"
+      }
+    }
+  ],
+  "version" : 3
+}