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

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")

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
+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, *)
+//#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

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

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: nonisolated(nonsending) () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A
+// async function pointer to Swift.withTaskCancellationHandler<A>(operation: nonisolated(nonsending) () 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

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: nonisolated(nonsending) () async throws -> A, onCancel: @Sendable () -> ()) async throws -> A
+// async function pointer to Swift.withTaskCancellationHandler<A>(operation: nonisolated(nonsending) () 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

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,19 @@ 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'
 
+// Hashable for (Throwing)AsyncStream
+// These are just @usableFromInline:
+Var AsyncStream.Continuation.storage is a new API without '@available'
+Var AsyncThrowingStream.Continuation.storage is a new API without '@available'
+Class AsyncStream._Storage is a new API without '@available'
+Class AsyncThrowingStream._Storage is a new API without '@available'
+
+// 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: nonisolated(nonsending) () 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.)