Merge pull request swiftlang#36590 from jckarter/task-cancellation-handler-abi

Fix calling convention for withCompletionHandler runtime calls.

Author: DougGregor

include/swift/ABI/TaskStatus.h

@@ -254,7 +254,7 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
 /// subsequently used.
 class CancellationNotificationStatusRecord : public TaskStatusRecord {
 public:
-  using FunctionType = void (void *);
+  using FunctionType = SWIFT_CC(swift) void (SWIFT_CONTEXT void *);
 
 private:
   FunctionType * __ptrauth_swift_cancellation_notification_function Function;

include/swift/Runtime/Concurrency.h

@@ -338,7 +338,8 @@ bool swift_task_isCancelled(AsyncTask* task);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 CancellationNotificationStatusRecord*
 swift_task_addCancellationHandler(
-    CancellationNotificationStatusRecord::FunctionType handler);
+    CancellationNotificationStatusRecord::FunctionType handler,
+    void *handlerContext);
 
 /// Remove the passed cancellation record from the task.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -132,8 +132,9 @@ OVERRIDE_TASK(continuation_throwingResumeWithError, void,
 OVERRIDE_TASK(task_addCancellationHandler,
               CancellationNotificationStatusRecord *,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
-              (CancellationNotificationStatusRecord::FunctionType handler),
-              (handler))
+              (CancellationNotificationStatusRecord::FunctionType handler,
+               void *context),
+              (handler, context))
 
 OVERRIDE_TASK(task_removeCancellationHandler, void,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,

stdlib/public/Concurrency/Task.cpp

@@ -789,12 +789,13 @@ bool swift::swift_task_isCancelled(AsyncTask *task) {
 SWIFT_CC(swift)
 static CancellationNotificationStatusRecord*
 swift_task_addCancellationHandlerImpl(
-    CancellationNotificationStatusRecord::FunctionType handler) {
+    CancellationNotificationStatusRecord::FunctionType handler,
+    void *context) {
   void *allocation =
       swift_task_alloc(sizeof(CancellationNotificationStatusRecord));
   auto *record =
       new (allocation) CancellationNotificationStatusRecord(
-          handler, /*arg=*/nullptr);
+          handler, context);
 
   swift_task_addStatusRecord(record);
   return record;

test/Concurrency/Runtime/cancellation_handler.swift

@@ -0,0 +1,21 @@
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency %import-libdispatch)
+
+class Canary {
+  deinit {
+    print("canary died")
+  }
+}
+
+do {
+  let task = Task.runDetached {
+    let canary = Canary()
+    _ = await Task.withCancellationHandler {
+      print(canary)
+    } operation: {
+      await Task.sleep(1_000_000)
+    }
+  }
+  task.cancel()
+}
+// CHECK: Canary
+// CHECK-NEXT: canary died