[Concurrency] Fix missing inherit executor on withTaskCancellationHandler (swiftlang#62546)

Author: ktoso

stdlib/public/BackDeployConcurrency/TaskCancellation.swift

@@ -26,8 +26,19 @@ import Swift
 ///
 /// Doesn't check for cancellation, and always executes the passed `operation`.
 ///
-/// This function returns immediately and never suspends.
+/// The `operation` executes on the calling execution context and does not suspend by itself,
+/// unless the code contained within the closure does. If cancellation occurs while the
+/// operation is running, the cancellation `handler` will execute *concurrently* with the `operation`.
+///
+/// ### Already cancelled tasks
+/// When `withTaskCancellationHandler` is used in a `Task` that has already been cancelled,
+/// the `onCancel` cancellation ``handler`` will be executed immediately before operation gets
+/// to execute. This allows the cancellation handler to set some external "cancelled" flag that the
+/// operation may be *atomically* checking for in order to avoid performing any actual work once
+/// the operation gets to run.
+@_unsafeInheritExecutor // the operation runs on the same executor as we start out with
 @available(SwiftStdlib 5.1, *)
+@_backDeploy(before: SwiftStdlib 5.8)
 public func withTaskCancellationHandler<T>(
   operation: () async throws -> T,
   onCancel handler: @Sendable () -> Void
@@ -92,10 +103,12 @@ public struct CancellationError: Error {
   public init() {}
 }
 
+@usableFromInline
 @available(SwiftStdlib 5.1, *)
 @_silgen_name("swift_task_addCancellationHandler")
 func _taskAddCancellationHandler(handler: () -> Void) -> UnsafeRawPointer /*CancellationNotificationStatusRecord*/
 
+@usableFromInline
 @available(SwiftStdlib 5.1, *)
 @_silgen_name("swift_task_removeCancellationHandler")
 func _taskRemoveCancellationHandler(

stdlib/public/Concurrency/TaskCancellation.swift

@@ -26,8 +26,19 @@ import Swift
 ///
 /// Doesn't check for cancellation, and always executes the passed `operation`.
 ///
-/// This function returns immediately and never suspends.
+/// The `operation` executes on the calling execution context and does not suspend by itself,
+/// unless the code contained within the closure does. If cancellation occurs while the
+/// operation is running, the cancellation `handler` will execute *concurrently* with the `operation`.
+///
+/// ### Already cancelled tasks
+/// When `withTaskCancellationHandler` is used in a `Task` that has already been cancelled,
+/// the `onCancel` cancellation ``handler`` will be executed immediately before operation gets
+/// to execute. This allows the cancellation handler to set some external "cancelled" flag that the
+/// operation may be *atomically* checking for in order to avoid performing any actual work once
+/// the operation gets to run.
+@_unsafeInheritExecutor // the operation runs on the same executor as we start out with
 @available(SwiftStdlib 5.1, *)
+@_backDeploy(before: SwiftStdlib 5.8)
 public func withTaskCancellationHandler<T>(
   operation: () async throws -> T,
   onCancel handler: @Sendable () -> Void
@@ -92,10 +103,12 @@ public struct CancellationError: Error {
   public init() {}
 }
 
+@usableFromInline
 @available(SwiftStdlib 5.1, *)
 @_silgen_name("swift_task_addCancellationHandler")
 func _taskAddCancellationHandler(handler: () -> Void) -> UnsafeRawPointer /*CancellationNotificationStatusRecord*/
 
+@usableFromInline
 @available(SwiftStdlib 5.1, *)
 @_silgen_name("swift_task_removeCancellationHandler")
 func _taskRemoveCancellationHandler(

test/Concurrency/Runtime/async_task_cancellation_taskGroup.swift

@@ -0,0 +1,40 @@
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking %import-libdispatch -parse-as-library) | %FileCheck %s --dump-input=always
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: libdispatch
+
+// rdar://76038845
+// REQUIRES: concurrency_runtime
+
+import Dispatch
+
+@available(SwiftStdlib 5.1, *)
+func test_detach_cancel_taskGroup() async {
+  print(#function) // CHECK: test_detach_cancel_taskGroup
+
+  await withTaskGroup(of: Void.self) { group in
+    group.cancelAll() // immediately cancel the group
+    print("group.cancel()") // CHECK: group.cancel()
+
+    group.addTask {
+      // immediately cancelled child task...
+      await withTaskCancellationHandler {
+        print("child: operation, was cancelled: \(Task.isCancelled)")
+      } onCancel: {
+        print("child: onCancel, was cancelled: \(Task.isCancelled)")
+      }
+    }
+    // CHECK: child: onCancel, was cancelled: true
+    // CHECK: child: operation, was cancelled: true
+  }
+
+  print("done") // CHECK: done
+}
+
+@available(SwiftStdlib 5.1, *)
+@main struct Main {
+  static func main() async {
+    await test_detach_cancel_taskGroup()
+  }
+}

test/Concurrency/actor_withCancellationHandler.swift

@@ -0,0 +1,20 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -emit-module-path %t/OtherActors.swiftmodule -module-name OtherActors %S/Inputs/OtherActors.swift -disable-availability-checking
+// RUN: %target-typecheck-verify-swift -I %t  -disable-availability-checking -warn-concurrency -parse-as-library
+// REQUIRES: concurrency
+
+actor foo {
+  var t: Task<Void, Error>?
+
+  func access() {}
+
+  func bar() {
+    self.t = Task {
+      await withTaskCancellationHandler {
+        self.access()
+      } onCancel: { @Sendable in
+
+      }
+    }
+  }
+}