[Concurrency] Task priority escalation handler API

Author: ktoso

include/swift/ABI/TaskStatus.h

@@ -242,7 +242,9 @@ class CancellationNotificationStatusRecord : public TaskStatusRecord {
       : TaskStatusRecord(TaskStatusRecordKind::CancellationNotification),
         Function(fn), Argument(arg) {}
 
-  void run() { Function(Argument); }
+  void run() {
+    Function(Argument);
+  }
 
   static bool classof(const TaskStatusRecord *record) {
     return record->getKind() == TaskStatusRecordKind::CancellationNotification;
@@ -259,7 +261,7 @@ class CancellationNotificationStatusRecord : public TaskStatusRecord {
 /// subsequently used.
 class EscalationNotificationStatusRecord : public TaskStatusRecord {
 public:
-  using FunctionType = void(void *, JobPriority);
+  using FunctionType = SWIFT_CC(swift) void(JobPriority, SWIFT_CONTEXT void *);
 
 private:
   FunctionType *__ptrauth_swift_escalation_notification_function Function;
@@ -268,9 +270,12 @@ class EscalationNotificationStatusRecord : public TaskStatusRecord {
 public:
   EscalationNotificationStatusRecord(FunctionType *fn, void *arg)
       : TaskStatusRecord(TaskStatusRecordKind::EscalationNotification),
-        Function(fn), Argument(arg) {}
+        Function(fn), Argument(arg) {
+  }
 
-  void run(JobPriority newPriority) { Function(Argument, newPriority); }
+  void run(JobPriority newPriority) {
+    Function(newPriority, Argument);
+  }
 
   static bool classof(const TaskStatusRecord *record) {
     return record->getKind() == TaskStatusRecordKind::EscalationNotification;

include/swift/Runtime/Concurrency.h

@@ -609,6 +609,18 @@ SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_removeCancellationHandler(
     CancellationNotificationStatusRecord *record);
 
+/// Create and add an escalation record to the task.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+EscalationNotificationStatusRecord*
+swift_task_addEscalationHandler(
+    EscalationNotificationStatusRecord::FunctionType handler,
+    void *handlerContext);
+
+/// Remove the passed cancellation record from the task.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_removeEscalationHandler(
+    EscalationNotificationStatusRecord *record);
+
 /// Create a NullaryContinuationJob from a continuation.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 NullaryContinuationJob*

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -206,6 +206,17 @@ OVERRIDE_TASK(task_removeCancellationHandler, void,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
               (CancellationNotificationStatusRecord *record), (record))
 
+OVERRIDE_TASK(task_addEscalationHandler,
+              EscalationNotificationStatusRecord *,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
+              (EscalationNotificationStatusRecord::FunctionType handler,
+               void *context),
+              (handler, context))
+
+OVERRIDE_TASK(task_removeEscalationHandler, void,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
+              (EscalationNotificationStatusRecord *record), (record))
+
 OVERRIDE_TASK(task_createNullaryContinuationJob, NullaryContinuationJob *,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
               (size_t priority,

stdlib/public/Concurrency/Task.cpp

@@ -1766,6 +1766,31 @@ static void swift_task_removeCancellationHandlerImpl(
   swift_task_dealloc(record);
 }
 
+SWIFT_CC(swift)
+static EscalationNotificationStatusRecord*
+swift_task_addEscalationHandlerImpl(
+    EscalationNotificationStatusRecord::FunctionType handler,
+    void *context) {
+  void *allocation =
+      swift_task_alloc(sizeof(EscalationNotificationStatusRecord));
+  auto unsigned_handler = swift_auth_code(handler, 3848); // FIXME: fix this number for ptrauth
+  auto *record = ::new (allocation)
+      EscalationNotificationStatusRecord(handler, context);
+
+  addStatusRecordToSelf(record, [&](ActiveTaskStatus oldStatus, ActiveTaskStatus& newStatus) {
+    return true;
+  });
+
+  return record;
+}
+
+SWIFT_CC(swift)
+static void swift_task_removeEscalationHandlerImpl(
+    EscalationNotificationStatusRecord *record) {
+  removeStatusRecordFromSelf(record);
+  swift_task_dealloc(record);
+}
+
 SWIFT_CC(swift)
 static NullaryContinuationJob*
 swift_task_createNullaryContinuationJobImpl(

stdlib/public/Concurrency/Task.swift

@@ -929,7 +929,7 @@ extension Task {
   ///   - task: the task which to escalate the priority of
   ///   - newPriority: the new priority the task should continue executing on
   @available(SwiftStdlib 9999, *)
-  public static func escalatePriority(_ task: UnsafeCurrentTask, to newPriority: TaskPriority) {
+  public static func escalatePriority(_ task: Task, to newPriority: TaskPriority) {
     _taskEscalate(task._task, newPriority: newPriority.rawValue)
   }
 
@@ -964,6 +964,56 @@ extension Task {
   }
 }
 
+/// Runs the passed `operation` while registering a task priority escalation handler.
+/// The handler will be triggered concurrently to the current task if the current
+/// is subject to priority escalation.
+///
+/// The handler may perform additional actions upon priority escalation,
+/// but cannot influence how the escalation influences the task, i.e. the task's
+/// priority will be escalated regardless of actions performed in the handler.
+///
+/// If multiple task escalation handlers are nester they will all be triggered.
+///
+/// Task escalation propagates through structured concurrency child-tasks.
+///
+/// - Parameters:
+///   - operation: the operation during which to listen for priority escalation
+///   - handler: handler to invoke, concurrently to `operation`,
+///              when priority escalation happens
+/// - Returns: the value returned by `operation`
+/// - Throws: when the `operation` throws an error
+@available(SwiftStdlib 9999, *)
+public func withTaskPriorityEscalationHandler<T, E>(
+  operation: () async throws(E) -> T,
+  onEscalate handler: @Sendable (TaskPriority) -> Void,
+  isolation: isolated (any Actor)? = #isolation
+) async throws(E) -> T {
+  // NOTE: We have to create the closure beforehand as otherwise it seems
+  // the task-local allocator may be used and we end up violating stack-discipline
+  // when releasing the handler closure vs. the record.
+  let handler0: (UInt8) -> Void = {
+    handler(TaskPriority(rawValue: $0))
+  }
+  let record = _taskAddEscalationHandler(handler: handler0)
+  defer { _taskRemoveEscalationHandler(record: record) }
+
+  return try await operation()
+}
+
+@usableFromInline
+@available(SwiftStdlib 9999, *)
+@_silgen_name("swift_task_addEscalationHandler")
+func _taskAddEscalationHandler(
+  handler: (UInt8) -> Void
+) -> UnsafeRawPointer /*EscalationNotificationStatusRecord*/
+
+@usableFromInline
+@available(SwiftStdlib 9999, *)
+@_silgen_name("swift_task_removeEscalationHandler")
+func _taskRemoveEscalationHandler(
+  record: UnsafeRawPointer /*EscalationNotificationStatusRecord*/
+)
+
 // ==== UnsafeCurrentTask ------------------------------------------------------
 
 /// Calls a closure with an unsafe reference to the current task.

test/Concurrency/async_task_escalate_priority.swift

@@ -26,8 +26,13 @@ import Darwin
 @preconcurrency import Dispatch
 import StdlibUnittest
 
+func TESTTEST_taskAddCancellationHandler(handler: () -> Void) {
+  handler()
+}
+func TESTTEST_taskAddEscalationHandler(handler: (TaskPriority) -> Void, t: TaskPriority) { handler(t) }
+
 func loopUntil(priority: TaskPriority) async {
-  var loops = 100
+  var loops = 10
   var currentPriority = Task.currentPriority
   while (currentPriority != priority) {
     print("Current priority = \(currentPriority) != \(priority)")
@@ -99,6 +104,43 @@ func testNestedTaskPriority(basePri: TaskPriority, curPri: TaskPriority) async {
           Task.escalatePriority(task, to: .default)
           sem2.wait()
         }
+
+        tests.test("Trigger task escalation handler") {
+          let sem1 = DispatchSemaphore(value: 0)
+          let sem2 = DispatchSemaphore(value: 0)
+          let semEscalated = DispatchSemaphore(value: 0)
+
+          let task = Task(priority: .background) {
+            let _ = expectedBasePri(priority: .background)
+
+            await withTaskPriorityEscalationHandler {
+              print("in withTaskPriorityEscalationHandler, Task.currentPriority = \(Task.currentPriority)")
+
+              // Wait until task is running before asking to be escalated
+              sem1.signal()
+              sleep(1)
+
+              await loopUntil(priority: .default)
+              print("in withTaskPriorityEscalationHandler, after loop, Task.currentPriority = \(Task.currentPriority)")
+            } onEscalate: { newPriority in
+              print("in onEscalate Task.currentPriority = \(Task.currentPriority)")
+              print("in onEscalate newPriority = \(newPriority)")
+              precondition(newPriority == .default)
+              semEscalated.signal()
+            }
+
+            print("Current priority = \(Task.currentPriority)")
+            print("after withTaskPriorityEscalationHandler")
+            sem2.signal()
+          }
+
+          // Wait till child runs and asks to be escalated
+          sem1.wait()
+          task.cancel() // just checking the records don't stomp onto each other somehow
+          Task.escalatePriority(task, to: .default)
+          semEscalated.wait()
+          sem2.wait()
+        }
       }
 
       await runAllTestsAsync()