[Concurrency] Set thread base priority when running escalated Tasks

Author: Bryce-MW

include/swift/ABI/Task.h

@@ -29,6 +29,13 @@
 #include "bitset"
 #include "queue" // TODO: remove and replace with our own mpsc
 
+// Does the runtime integrate with libdispatch?
+#if defined(SWIFT_CONCURRENCY_USES_DISPATCH)
+#define SWIFT_CONCURRENCY_ENABLE_DISPATCH SWIFT_CONCURRENCY_USES_DISPATCH
+#else
+#define SWIFT_CONCURRENCY_ENABLE_DISPATCH 0
+#endif
+
 // Does the runtime provide priority escalation support?
 #ifndef SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
 #if SWIFT_CONCURRENCY_ENABLE_DISPATCH && \
@@ -422,7 +429,11 @@ class AsyncTask : public Job {
   ///
   /// Generally this should be done immediately after updating
   /// ActiveTask.
-  void flagAsRunning();
+  ///
+  /// If the return value is non-zero, it must be passed to
+  /// swift_dispatch_thread_reset_override_self. If it is zero,
+  /// calling that function is optional.
+  uint32_t flagAsRunning();
 
   /// Flag that this task is now suspended with information about what it is
   /// waiting on.

include/swift/Runtime/Concurrency.h

@@ -38,13 +38,6 @@
 #define SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL 0
 #endif
 
-// Does the runtime integrate with libdispatch?
-#if defined(SWIFT_CONCURRENCY_USES_DISPATCH)
-#define SWIFT_CONCURRENCY_ENABLE_DISPATCH SWIFT_CONCURRENCY_USES_DISPATCH
-#else
-#define SWIFT_CONCURRENCY_ENABLE_DISPATCH 0
-#endif
-
 namespace swift {
 class DefaultActor;
 class TaskOptionRecord;

include/swift/Runtime/DispatchShims.h

@@ -48,6 +48,27 @@ swift_dispatch_thread_override_self(qos_class_t override_qos) {
   return 0;
 }
 
+static inline uint32_t
+swift_dispatch_thread_override_self_with_base(qos_class_t override_qos, qos_class_t base_qos) {
+
+  if (__builtin_available(macOS 27.0, iOS 27.0, tvOS 27.0, watchOS 27.0, *)) {
+    return dispatch_thread_override_self_with_base(override_qos, base_qos);
+  } else if (__builtin_available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)) {
+    // If we don't have the ability to set our base qos correctly, at least set the override
+    // We want to return 0 here because we have nothing to reset in this case
+    (void) dispatch_thread_override_self(override_qos);
+  }
+
+  return 0;
+}
+
+static inline void
+swift_dispatch_thread_reset_override_self(uint32_t opaque) {
+  if (__builtin_available(macOS 27.0, iOS 27.0, tvOS 27.0, watchOS 27.0, *)) {
+    dispatch_thread_reset_override_self(opaque);
+  }
+}
+
 static inline int
 swift_dispatch_lock_override_start_with_debounce(dispatch_lock_t *lock_addr,
    dispatch_tid_t expected_thread, qos_class_t override_to_apply) {

stdlib/public/Concurrency/Actor.cpp

@@ -237,12 +237,14 @@ void swift::runJobInEstablishedExecutorContext(Job *job) {
     // current thread.  If the task suspends somewhere, it should
     // update the task status appropriately; we don't need to update
     // it afterwards.
-    task->flagAsRunning();
+    uint32_t dispatch_opaque_priority = task->flagAsRunning();
 
     auto traceHandle = concurrency::trace::job_run_begin(job);
     task->runInFullyEstablishedContext();
     concurrency::trace::job_run_end(traceHandle);
 
+    swift_dispatch_thread_reset_override_self(dispatch_opaque_priority);
+
     assert(ActiveTask::get() == nullptr &&
            "active task wasn't cleared before suspending?");
     if (oldTask) ActiveTask::set(oldTask);

stdlib/public/Concurrency/Task.cpp

@@ -132,7 +132,8 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
       SWIFT_TASK_DEBUG_LOG("task %p waiting on task %p, completed immediately",
                            waitingTask, this);
       _swift_tsan_acquire(static_cast<Job *>(this));
-      if (contextInitialized) waitingTask->flagAsRunning();
+      // Continuting execution inline so we shouldn't have to reset priority here
+      if (contextInitialized) assert(waitingTask->flagAsRunning() == 0);
       // The task is done; we don't need to wait.
       return queueHead.getStatus();
 
@@ -1666,8 +1667,9 @@ static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
   // we try to tail-call.
   } while (false);
 #else
-  // Restore the running state of the task and resume it.
-  task->flagAsRunning();
+  // Restore the running state of the task and resume it. We are on the same
+  // thread as before so dispatch shouldn't pass us anything to restore
+  assert(task->flagAsRunning() == 0);
 #endif /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
 
   if (context->isExecutorSwitchForced())

stdlib/public/Concurrency/TaskGroup.cpp

@@ -1822,7 +1822,8 @@ reevaluate_if_taskgroup_has_results:;
     // We're going back to running the task, so if we suspended before,
     // we need to flag it as running again.
     if (hasSuspended) {
-      waitingTask->flagAsRunning();
+      // Continuting execution inline so we shouldn't have to reset priority here
+      assert(waitingTask->flagAsRunning() == 0);
     }
 
     // Success! We are allowed to poll.

stdlib/public/Concurrency/TaskPrivate.h

@@ -935,32 +935,40 @@ inline bool AsyncTask::isCancelled() const {
                           .isCancelled();
 }
 
-inline void AsyncTask::flagAsRunning() {
+inline uint32_t AsyncTask::flagAsRunning() {
 
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
   dispatch_thread_override_info_s threadOverrideInfo;
   threadOverrideInfo = swift_dispatch_thread_get_current_override_qos_floor();
   qos_class_t overrideFloor = threadOverrideInfo.override_qos_floor;
+  qos_class_t basePriorityCeil = overrideFloor;
+  qos_class_t taskBasePriority = (qos_class_t) _private().BasePriority;
 #endif
 
   auto oldStatus = _private()._status().load(std::memory_order_relaxed);
   assert(!oldStatus.isRunning());
   assert(!oldStatus.isComplete());
 
+  uint32_t dispatch_opaque_priority = 0;
   if (!oldStatus.hasTaskDependency()) {
     SWIFT_TASK_DEBUG_LOG("%p->flagAsRunning() with no task dependency", this);
     assert(_private().dependencyRecord == nullptr);
 
     while (true) {
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
-      // Task's priority is greater than the thread's - do a self escalation
+      // If the base priority is not equal to the current override floor then
+      // dispqatch may need to apply the base priority to the thread. If the
+      // current priority is higher than the override floor, then dispatch may
+      // need to apply a self-override. In either case, call into dispatch to
+      // do this.
       qos_class_t maxTaskPriority = (qos_class_t) oldStatus.getStoredPriority();
-      if (threadOverrideInfo.can_override && (maxTaskPriority > overrideFloor)) {
-        SWIFT_TASK_DEBUG_LOG("[Override] Self-override thread with oq_floor %#x to match task %p's max priority %#x",
-            overrideFloor, this, maxTaskPriority);
+      if (threadOverrideInfo.can_override && (taskBasePriority != basePriorityCeil || maxTaskPriority > overrideFloor)) {
+        SWIFT_TASK_DEBUG_LOG("[Override] Self-override thread with oq_floor %#x to match task %p's max priority %#x and base priority %#x",
+                             overrideFloor, this, maxTaskPriority, taskBasePriority);
 
-        (void) swift_dispatch_thread_override_self(maxTaskPriority);
+        dispatch_opaque_priority = swift_dispatch_thread_override_self_with_base(maxTaskPriority, taskBasePriority);
         overrideFloor = maxTaskPriority;
+        basePriorityCeil = taskBasePriority;
       }
 #endif
       // Set self as executor and remove escalation bit if any - the task's
@@ -989,14 +997,19 @@ inline void AsyncTask::flagAsRunning() {
                        ActiveTaskStatus& newStatus) {
 
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
-      // Task's priority is greater than the thread's - do a self escalation
+      // If the base priority is not equal to the current override floor then
+      // dispqatch may need to apply the base priority to the thread. If the
+      // current priority is higher than the override floor, then dispatch may
+      // need to apply a self-override. In either case, call into dispatch to
+      // do this.
       qos_class_t maxTaskPriority = (qos_class_t) oldStatus.getStoredPriority();
-      if (threadOverrideInfo.can_override && (maxTaskPriority > overrideFloor)) {
-        SWIFT_TASK_DEBUG_LOG("[Override] Self-override thread with oq_floor %#x to match task %p's max priority %#x",
-            overrideFloor, this, maxTaskPriority);
+      if (threadOverrideInfo.can_override && (taskBasePriority != basePriorityCeil || maxTaskPriority > overrideFloor)) {
+        SWIFT_TASK_DEBUG_LOG("[Override] Self-override thread with oq_floor %#x to match task %p's max priority %#x and base priority %#x",
+                             overrideFloor, this, maxTaskPriority, taskBasePriority);
 
-        (void) swift_dispatch_thread_override_self(maxTaskPriority);
+        dispatch_opaque_priority = swift_dispatch_thread_override_self_with_base(maxTaskPriority, taskBasePriority);
         overrideFloor = maxTaskPriority;
+        basePriorityCeil = taskBasePriority;
       }
 #endif
       // Set self as executor and remove escalation bit if any - the task's
@@ -1012,7 +1025,7 @@ inline void AsyncTask::flagAsRunning() {
     swift_task_enterThreadLocalContext(
         (char *)&_private().ExclusivityAccessSet[0]);
   }
-
+  return dispatch_opaque_priority;
 }
 
 /// TODO (rokhinip): We need the handoff of the thread to the next executor to