[Concurrency] Adjust task escalation APIs to SE accepted shapes (#79988)

* [Concurrency] Adjust task escalation APIs to SE accepted shapes

* adjust test a little bit

* Fix closure lifetime in withTaskPriorityEscalationHandler

* avoid bringing workaround func into abi by marking AEIC

Author: ktoso

include/swift/ABI/MetadataValues.h

@@ -1740,7 +1740,7 @@ namespace SpecialPointerAuthDiscriminators {
   const uint16_t AsyncContextResume = 0xd707; // = 55047
   const uint16_t AsyncContextYield = 0xe207; // = 57863
   const uint16_t CancellationNotificationFunction = 0x1933; // = 6451
-  const uint16_t EscalationNotificationFunction = 0xf59d; // = 62877
+  const uint16_t EscalationNotificationFunction = 0x7861; // = 30817
   const uint16_t AsyncThinNullaryFunction = 0x0f08; // = 3848
   const uint16_t AsyncFutureFunction = 0x720f; // = 29199
 

include/swift/ABI/TaskStatus.h

@@ -261,7 +261,7 @@ class CancellationNotificationStatusRecord : public TaskStatusRecord {
 /// subsequently used.
 class EscalationNotificationStatusRecord : public TaskStatusRecord {
 public:
-  using FunctionType = SWIFT_CC(swift) void(JobPriority, SWIFT_CONTEXT void *);
+  using FunctionType = SWIFT_CC(swift) void(JobPriority, JobPriority, SWIFT_CONTEXT void *);
 
 private:
   FunctionType *__ptrauth_swift_escalation_notification_function Function;
@@ -273,8 +273,8 @@ class EscalationNotificationStatusRecord : public TaskStatusRecord {
         Function(fn), Argument(arg) {
   }
 
-  void run(JobPriority newPriority) {
-    Function(newPriority, Argument);
+  void run(JobPriority oldPriority, JobPriority newPriority) {
+    Function(oldPriority, newPriority, Argument);
   }
 
   static bool classof(const TaskStatusRecord *record) {
@@ -439,7 +439,8 @@ class TaskDependencyStatusRecord : public TaskStatusRecord {
     DependentOn.Executor = executor;
   }
 
-  void performEscalationAction(JobPriority newPriority);
+  void performEscalationAction(
+      JobPriority oldPriority, JobPriority newPriority);
 };
 
 } // end namespace swift

stdlib/public/Concurrency/Task+PriorityEscalation.swift

@@ -17,7 +17,7 @@ import Swift
 
 @available(SwiftStdlib 6.2, *)
 extension Task {
-  /// Escalate the task `priority` of the passed in task to the `newPriority`.
+  /// Manually escalate the task `priority` of this task to the `newPriority`.
   ///
   /// - Warning: This API should rarely be used, and instead you can rely on
   ///   structured concurrency and implicit priority escalation which happens
@@ -43,8 +43,8 @@ extension Task {
   ///   - task: the task which to escalate the priority of
   ///   - newPriority: the new priority the task should continue executing on
   @available(SwiftStdlib 6.2, *)
-  public static func escalatePriority(_ task: Task, to newPriority: TaskPriority) {
-    _taskEscalate(task._task, newPriority: newPriority.rawValue)
+  public func escalatePriority(to newPriority: TaskPriority) {
+    _taskEscalate(self._task, newPriority: newPriority.rawValue)
   }
 }
 
@@ -76,8 +76,8 @@ extension UnsafeCurrentTask {
   ///   - task: the task which to escalate the priority of
   ///   - newPriority: the new priority the task should continue executing on
   @available(SwiftStdlib 6.2, *)
-  public static func escalatePriority(_ task: UnsafeCurrentTask, to newPriority: TaskPriority) {
-    unsafe _taskEscalate(task._task, newPriority: newPriority.rawValue)
+  public func escalatePriority(to newPriority: TaskPriority) {
+    unsafe _taskEscalate(self._task, newPriority: newPriority.rawValue)
   }
 }
 
@@ -107,15 +107,24 @@ extension UnsafeCurrentTask {
 @available(SwiftStdlib 6.2, *)
 public func withTaskPriorityEscalationHandler<T, E>(
   operation: () async throws(E) -> T,
-  onPriorityEscalated handler: @Sendable (TaskPriority) -> Void,
+  onPriorityEscalated handler: @Sendable (TaskPriority, TaskPriority) -> Void,
+  isolation: isolated (any Actor)? = #isolation
+) async throws(E) -> T {
+  return try await __withTaskPriorityEscalationHandler0(
+    operation: operation,
+    onPriorityEscalated: {
+      handler(TaskPriority(rawValue: $0), TaskPriority(rawValue: $1))
+    })
+}
+
+// Method necessary in order to avoid the handler0 to be destroyed too eagerly.
+@available(SwiftStdlib 6.2, *)
+@_alwaysEmitIntoClient
+func __withTaskPriorityEscalationHandler0<T, E>(
+  operation: () async throws(E) -> T,
+  onPriorityEscalated handler0: @Sendable (UInt8, UInt8) -> 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 = unsafe _taskAddPriorityEscalationHandler(handler: handler0)
   defer { unsafe _taskRemovePriorityEscalationHandler(record: record) }
 

stdlib/public/Concurrency/Task.swift

@@ -1430,7 +1430,7 @@ func _enqueueJobGlobalWithDeadline(_ seconds: Int64, _ nanoseconds: Int64,
 @available(SwiftStdlib 6.2, *)
 @_silgen_name("swift_task_addPriorityEscalationHandler")
 func _taskAddPriorityEscalationHandler(
-  handler: (UInt8) -> Void
+  handler: (UInt8, UInt8) -> Void
 ) -> UnsafeRawPointer /*EscalationNotificationStatusRecord*/
 
 @usableFromInline

stdlib/public/Concurrency/TaskPrivate.h

@@ -1125,7 +1125,9 @@ void AsyncTask::flagAsSuspended(TaskDependencyStatusRecord *dependencyStatusReco
     // reference to the dependencyStatusRecord or its contents, once we have
     // published it in the ActiveTaskStatus since someone else could
     // concurrently made us runnable.
-    dependencyStatusRecord->performEscalationAction(newStatus.getStoredPriority());
+    dependencyStatusRecord->performEscalationAction(
+        oldStatus.getStoredPriority(),
+        newStatus.getStoredPriority());
 
     // Always add the dependency status record
     return true;

stdlib/public/Concurrency/TaskStatus.cpp

@@ -1047,6 +1047,7 @@ static void swift_task_cancelImpl(AsyncTask *task) {
 
 /// Perform any escalation actions required by the given record.
 static void performEscalationAction(TaskStatusRecord *record,
+                                    JobPriority oldPriority,
                                     JobPriority newPriority) {
   switch (record->getKind()) {
   // Child tasks need to be recursively escalated.
@@ -1067,15 +1068,17 @@ static void performEscalationAction(TaskStatusRecord *record,
   case TaskStatusRecordKind::EscalationNotification:  {
     auto notification =
       cast<EscalationNotificationStatusRecord>(record);
-    notification->run(newPriority);
+    SWIFT_TASK_DEBUG_LOG("[Dependency] Trigger task escalation handler record %p, escalate from %#x to %#x",
+                         record, oldPriority, newPriority);
+    notification->run(oldPriority, newPriority);
     return;
   }
 
   case TaskStatusRecordKind::TaskDependency: {
     auto dependencyRecord = cast<TaskDependencyStatusRecord>(record);
-    SWIFT_TASK_DEBUG_LOG("[Dependency] Escalating a task dependency record %p to %#x",
-                    record, newPriority);
-    dependencyRecord->performEscalationAction(newPriority);
+    SWIFT_TASK_DEBUG_LOG("[Dependency] Escalating a task dependency record %p from %#x to %#x",
+                    record, oldPriority, newPriority);
+    dependencyRecord->performEscalationAction(oldPriority, newPriority);
     return;
   }
 
@@ -1101,17 +1104,17 @@ static void performEscalationAction(TaskStatusRecord *record,
 SWIFT_CC(swift)
 JobPriority
 static swift_task_escalateImpl(AsyncTask *task, JobPriority newPriority) {
-
   SWIFT_TASK_DEBUG_LOG("Escalating %p to %#zx priority", task, newPriority);
   auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
+  auto oldPriority = oldStatus.getStoredPriority();
   auto newStatus = oldStatus;
 
   while (true) {
     // Fast path: check that the stored priority is already at least
     // as high as the desired priority.
-    if (oldStatus.getStoredPriority() >= newPriority) {
-      SWIFT_TASK_DEBUG_LOG("Task is already at %#zx priority", oldStatus.getStoredPriority());
-      return oldStatus.getStoredPriority();
+    if (oldPriority >= newPriority) {
+      SWIFT_TASK_DEBUG_LOG("Task is already at %#zx priority", oldPriority);
+      return oldPriority;
     }
 
     if (oldStatus.isRunning() || oldStatus.isEnqueued()) {
@@ -1145,8 +1148,11 @@ static swift_task_escalateImpl(AsyncTask *task, JobPriority newPriority) {
     taskStatus = (ActiveTaskStatus *) &task->_private()._status();
     executionLock = (dispatch_lock_t *) ((char*)taskStatus + ActiveTaskStatus::executionLockOffset());
 
-    SWIFT_TASK_DEBUG_LOG("[Override] Escalating %p which is running on %#x to %#x", task, newStatus.currentExecutionLockOwner(), newPriority);
-    swift_dispatch_lock_override_start_with_debounce(executionLock, newStatus.currentExecutionLockOwner(), (qos_class_t) newPriority);
+    SWIFT_TASK_DEBUG_LOG("[Override] Escalating %p which is running on %#x from %#x to %#x",
+                         task, newStatus.currentExecutionLockOwner(),
+                         oldPriority, newPriority);
+    swift_dispatch_lock_override_start_with_debounce(
+        executionLock, newStatus.currentExecutionLockOwner(), (qos_class_t) newPriority);
 #endif
   } else if (newStatus.isEnqueued()) {
     //  Task is not running, it's enqueued somewhere waiting to be run
@@ -1165,22 +1171,24 @@ static swift_task_escalateImpl(AsyncTask *task, JobPriority newPriority) {
     return newStatus.getStoredPriority();
   }
 
-  SWIFT_TASK_DEBUG_LOG("[Override] Escalating %p which is suspended to %#x", task, newPriority);
+  SWIFT_TASK_DEBUG_LOG("[Override] Escalating %p which is suspended from %#x to %#x",
+                       task, oldPriority, newPriority);
   // We must have at least one record - the task dependency one.
   assert(newStatus.getInnermostRecord() != NULL);
 
   withStatusRecordLock(task, newStatus, [&](ActiveTaskStatus status) {
     // We know that none of the escalation actions will recursively
     // modify the task status record list by adding or removing task records
     for (auto cur: status.records()) {
-      performEscalationAction(cur, newPriority);
+      performEscalationAction(cur, oldPriority, newPriority);
     }
   });
 
   return newStatus.getStoredPriority();
 }
 
-void TaskDependencyStatusRecord::performEscalationAction(JobPriority newPriority) {
+void TaskDependencyStatusRecord::performEscalationAction(
+    JobPriority oldPriority, JobPriority newPriority) {
   switch (this->DependencyKind) {
     case WaitingOnTask:
       SWIFT_TASK_DEBUG_LOG("[Dependency] Escalate dependent task %p noted in %p record",
@@ -1202,8 +1210,8 @@ void TaskDependencyStatusRecord::performEscalationAction(JobPriority newPriority
         this->DependentOn.TaskGroup, this);
       break;
     case EnqueuedOnExecutor:
-      SWIFT_TASK_DEBUG_LOG("[Dependency] Escalate dependent executor %p noted in %p record",
-        this->DependentOn.Executor, this);
+      SWIFT_TASK_DEBUG_LOG("[Dependency] Escalate dependent executor %p noted in %p record from %#x to %#x",
+        this->DependentOn.Executor, this, oldPriority, newPriority);
       swift_executor_escalate(this->DependentOn.Executor, this->WaitingTask, newPriority);
       break;
   }