[Concurrency] implement withCancellationHandler via records

include/swift/Runtime/Concurrency.h

@@ -225,6 +225,17 @@ size_t swift_task_getJobFlags(AsyncTask* task);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 bool swift_task_isCancelled(AsyncTask* task);
 
+/// Create and add an cancellation record to the task.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+CancellationNotificationStatusRecord*
+swift_task_addCancellationHandler(
+    AsyncTask *task, CancellationNotificationStatusRecord::FunctionType handler);
+
+/// Remove the passed cancellation record from the task.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_removeCancellationHandler(
+    AsyncTask *task, CancellationNotificationStatusRecord *record);
+
 using TaskLocalValuesFragment = AsyncTask::TaskLocalValuesFragment;
 
 /// Get a task local value from the passed in task. Its Swift signature is

stdlib/public/Concurrency/Task.cpp

@@ -569,6 +569,25 @@ bool swift::swift_task_isCancelled(AsyncTask *task) {
   return task->isCancelled();
 }
 
+CancellationNotificationStatusRecord*
+swift::swift_task_addCancellationHandler(
+    AsyncTask *task, CancellationNotificationStatusRecord::FunctionType handler) {
+  void *allocation =
+      swift_task_alloc(task, sizeof(CancellationNotificationStatusRecord));
+  auto *record =
+      new (allocation) CancellationNotificationStatusRecord(
+          handler, /*arg=*/nullptr);
+
+  swift_task_addStatusRecord(task, record);
+  return record;
+}
+
+void swift::swift_task_removeCancellationHandler(
+    AsyncTask *task, CancellationNotificationStatusRecord *record) {
+  swift_task_removeStatusRecord(task, record);
+  swift_task_dealloc(task, record);
+}
+
 SWIFT_CC(swift)
 void swift::swift_continuation_logFailedCheck(const char *message) {
   swift_reportError(0, message);

stdlib/public/Concurrency/Task.swift

@@ -585,9 +585,6 @@ func getJobFlags(_ task: Builtin.NativeObject) -> Task.JobFlags
 @usableFromInline
 func _enqueueJobGlobal(_ task: Builtin.Job)
 
-@_silgen_name("swift_task_isCancelled")
-func isTaskCancelled(_ task: Builtin.NativeObject) -> Bool
-
 @available(*, deprecated)
 @_silgen_name("swift_task_runAndBlockThread")
 public func runAsyncAndBlock(_ asyncFun: @escaping () async -> ())

stdlib/public/Concurrency/TaskCancellation.swift

@@ -65,14 +65,23 @@ extension Task {
   ///
   /// Does not check for cancellation, and always executes the passed `operation`.
   ///
-  /// ### Suspension
   /// This function returns instantly and will never suspend.
-  /* @instantaneous */
   public static func withCancellationHandler<T>(
     handler: @concurrent () -> (),
     operation: () async throws -> T
-  ) async throws -> T {
-      fatalError("\(#function) not implemented yet.")
+  ) async rethrows -> T {
+    let task = Builtin.getCurrentAsyncTask()
+
+    guard !_taskIsCancelled(task) else {
+      // If the current task is already cancelled, run the handler immediately.
+      handler()
+      return try await operation()
+    }
+
+    let record = _taskAddCancellationHandler(task: task, handler: handler)
+    defer { _taskRemoveCancellationHandler(task: task, record: record) }
+
+    return try await operation()
   }
 
   /// The default cancellation thrown when a task is cancelled.
@@ -85,3 +94,15 @@ extension Task {
   }
 
 }
+
+@_silgen_name("swift_task_addCancellationHandler")
+func _taskAddCancellationHandler(
+  task: Builtin.NativeObject,
+  handler: @concurrent () -> ()
+) -> UnsafeRawPointer /*CancellationNotificationStatusRecord*/
+
+@_silgen_name("swift_task_removeCancellationHandler")
+func _taskRemoveCancellationHandler(
+  task: Builtin.NativeObject,
+  record: UnsafeRawPointer /*CancellationNotificationStatusRecord*/
+)

test/Concurrency/async_cancellation.swift

@@ -26,7 +26,7 @@ func test_cancellation_withCancellationHandler(_ anything: Any) async -> Picture
   let handle: Task.Handle<PictureData, Error> = Task.runDetached {
     let file = SomeFile()
 
-    return try await Task.withCancellationHandler(
+    return await Task.withCancellationHandler(
       handler: { file.close() },
       operation: {
       await test_cancellation_guard_isCancelled(file)