Task priority escalation on Apple platforms

A task can be in one of 4 states over its lifetime:

    (a) suspended
    (b) enqueued
    (c) running
    (d) completed

This change provides priority inversion avoidance support if a task gets
escalated when it is in state (a), (c), (d).

Radar-Id: rdar://problem/76127624

Author: rokhinip

include/swift/ABI/Task.h

@@ -290,6 +290,22 @@ class AsyncTask : public Job {
     return ResumeTask(ResumeContext); // 'return' forces tail call
   }
 
+  /// A task can have the following states:
+  ///   * suspended: In this state, a task is considered not runnable
+  ///   * enqueued: In this state, a task is considered runnable
+  ///   * running on a thread
+  ///   * completed
+  ///
+  /// The following state transitions are possible:
+  ///       suspended -> enqueued
+  ///       suspended -> running
+  ///       enqueued -> running
+  ///       running -> suspended
+  ///       running -> completed
+  ///       running -> enqueued
+  ///
+  /// The 4 methods below are how a task switches from one state to another.
+
   /// Flag that this task is now running.  This can update
   /// the priority stored in the job flags if the priority has been
   /// escalated.
@@ -298,14 +314,12 @@ class AsyncTask : public Job {
   /// ActiveTask.
   void flagAsRunning();
 
-  /// Flag that this task is now suspended.  This can update the
-  /// priority stored in the job flags if the priority hsa been
-  /// escalated.  Generally this should be done immediately after
-  /// clearing ActiveTask and immediately before enqueuing the task
-  /// somewhere.  TODO: record where the task is enqueued if
-  /// possible.
+  /// Flag that this task is now suspended.
   void flagAsSuspended();
 
+  /// Flag that the task is to be enqueued on the provided executor
+  void flagAsEnqueuedOnExecutor(ExecutorRef newExecutor);
+
   /// Flag that this task is now completed. This normally does not do anything
   /// but can be used to locally insert logging.
   void flagAsCompleted();

include/swift/Runtime/Concurrency.h

@@ -627,6 +627,17 @@ void swift_task_switch(SWIFT_ASYNC_CONTEXT AsyncContext *resumeToContext,
                        TaskContinuationFunction *resumeFunction,
                        ExecutorRef newExecutor);
 
+/// Mark a task for enqueue on a new executor and then enqueue it.
+///
+/// The resumption function pointer and continuation should be set
+/// appropriately in the task.
+///
+/// Generally you should call swift_task_switch to switch execution
+/// synchronously when possible.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void
+swift_task_enqueueTaskOnExecutor(AsyncTask *task, ExecutorRef executor);
+
 /// Enqueue the given job to run asynchronously on the given executor.
 ///
 /// The resumption function pointer and continuation should be set

include/swift/Runtime/DispatchShims.h

@@ -0,0 +1,64 @@
+//===--- DispatchShims.h - Shims for dispatch vended APIs --------------------*- C++ -*-//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_CONCURRENCY_DISPATCHSHIMS_H
+#define SWIFT_CONCURRENCY_DISPATCHSHIMS_H
+
+#include "Concurrency.h"
+
+#if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
+#include <dispatch/swift_concurrency_private.h>
+
+// Provide wrappers with runtime checks to make sure that the dispatch functions
+// are only called on OS-es where they are supported
+dispatch_thread_override_info_s
+swift_dispatch_thread_get_current_override_qos_floor()
+{
+  if (__builtin_available(macOS 9998, iOS 9998, tvOS 9998, watchOS 9998, *)) {
+    return dispatch_thread_get_current_override_qos_floor();
+  }
+
+  return (dispatch_thread_override_info_s) {0};
+}
+
+int
+swift_dispatch_thread_override_self(qos_class_t override_qos) {
+
+  if (__builtin_available(macOS 9998, iOS 9998, tvOS 9998, watchOS 9998, *)) {
+    return dispatch_thread_override_self(override_qos);
+  }
+
+  return 0;
+}
+
+int
+swift_dispatch_lock_override_start_with_debounce(dispatch_lock_t *lock_addr,
+   dispatch_tid_t expected_thread, qos_class_t override_to_apply) {
+
+  if (__builtin_available(macOS 9998, iOS 9998, tvOS 9998, watchOS 9998, *)) {
+    return dispatch_lock_override_start_with_debounce(lock_addr, expected_thread, override_to_apply);
+  }
+
+  return 0;
+}
+
+int
+swift_dispatch_lock_override_end(qos_class_t override_to_end) {
+  if (__builtin_available(macOS 9998, iOS 9998, tvOS 9998, watchOS 9998, *)) {
+    return dispatch_lock_override_end(override_to_end);
+  }
+
+  return 0;
+}
+#endif
+
+#endif /* SWIFT_CONCURRENCY_DISPATCHSHIMS_H */

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -159,6 +159,10 @@ OVERRIDE_TASK(task_suspend, AsyncTask *,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
               swift::, ,)
 
+OVERRIDE_TASK(task_enqueueTaskOnExecutor, void,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
+              (AsyncTask *task, ExecutorRef newExecutor), (task, newExecutor))
+
 OVERRIDE_TASK(continuation_init, AsyncTask *,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
               swift::, (ContinuationAsyncContext *context,

stdlib/public/Concurrency/Actor.cpp

@@ -1617,9 +1617,9 @@ static void swift_task_switchImpl(SWIFT_ASYNC_CONTEXT AsyncContext *resumeContex
   // executor.
   SWIFT_TASK_DEBUG_LOG("switch failed, task %p enqueued on executor %p", task,
                        newExecutor.getIdentity());
-  task->flagAsSuspended();
+
+  task->flagAsEnqueuedOnExecutor(newExecutor);
   _swift_task_clearCurrent();
-  swift_task_enqueue(task, newExecutor);
 }
 
 /*****************************************************************************/

stdlib/public/Concurrency/Task.cpp

@@ -145,9 +145,12 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
             queueHead, newQueueHead,
             /*success*/ std::memory_order_release,
             /*failure*/ std::memory_order_acquire)) {
+
       // Escalate the priority of this task based on the priority
       // of the waiting task.
-      swift_task_escalate(this, waitingTask->Flags.getPriority());
+      auto status = waitingTask->_private().Status.load(std::memory_order_relaxed);
+      swift_task_escalate(this, status.getStoredPriority());
+
       _swift_task_clearCurrent();
       return FutureFragment::Status::Executing;
     }
@@ -231,7 +234,7 @@ void AsyncTask::completeFuture(AsyncContext *context) {
 
     // Enqueue the waiter on the global executor.
     // TODO: allow waiters to fill in a suggested executor
-    swift_task_enqueueGlobal(waitingTask);
+    waitingTask->flagAsEnqueuedOnExecutor(ExecutorRef::generic());
 
     // Move to the next task.
     waitingTask = nextWaitingTask;
@@ -649,12 +652,7 @@ static AsyncTaskAndContext swift_task_create_commonImpl(
      basePriority = JobPriority::Default;
   }
 
-  SWIFT_TASK_DEBUG_LOG("Task's base priority = %d", basePriority);
-
-  // TODO (rokhinip): Figure out the semantics of the job priority and where
-  // it ought to be set conclusively - seems like it ought to be at enqueue
-  // time. For now, maintain current semantics of setting jobPriority as well.
-  jobFlags.setPriority(basePriority);
+  SWIFT_TASK_DEBUG_LOG("Task's base priority = %#x", basePriority);
 
   // Figure out the size of the header.
   size_t headerSize = sizeof(AsyncTask);
@@ -849,7 +847,7 @@ static AsyncTaskAndContext swift_task_create_commonImpl(
   // If we're supposed to enqueue the task, do so now.
   if (taskCreateFlags.enqueueJob()) {
     swift_retain(task);
-    swift_task_enqueue(task, executor);
+    task->flagAsEnqueuedOnExecutor(executor);
   }
 
   return {task, initialContext};
@@ -1018,6 +1016,13 @@ static AsyncTask *swift_task_suspendImpl() {
   return task;
 }
 
+SWIFT_CC(swift)
+static void
+swift_task_enqueueTaskOnExecutorImpl(AsyncTask *task, ExecutorRef executor)
+{
+  task->flagAsEnqueuedOnExecutor(executor);
+}
+
 SWIFT_CC(swift)
 static AsyncTask *swift_continuation_initImpl(ContinuationAsyncContext *context,
                                               AsyncContinuationFlags flags) {
@@ -1146,7 +1151,7 @@ static void resumeTaskAfterContinuation(AsyncTask *task,
   // to make a stronger best-effort attempt to catch racing attempts to
   // resume the continuation?
 
-  swift_task_enqueue(task, context->ResumeToExecutor);
+  task->flagAsEnqueuedOnExecutor(context->ResumeToExecutor);
 }
 
 SWIFT_CC(swift)

stdlib/public/Concurrency/TaskGroup.cpp

@@ -622,7 +622,7 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
         _swift_tsan_acquire(static_cast<Job *>(waitingTask));
 
         // TODO: allow the caller to suggest an executor
-        swift_task_enqueueGlobal(waitingTask);
+        waitingTask->flagAsEnqueuedOnExecutor(ExecutorRef::generic());
         return;
       } // else, try again
     }