Merge pull request #41088 from apple/rokhinip/76127624-task-priority-escalation

Task Priority Escalation

Author: rokhinip

include/swift/ABI/Task.h

@@ -290,23 +290,35 @@ 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.
   ///
   /// Generally this should be done immediately after updating
   /// ActiveTask.
   void flagAsRunning();
-  void flagAsRunning_slow();
-
-  /// 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();
-  void flagAsSuspended_slow();
+
+  /// 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.

include/swift/Reflection/ReflectionContext.h

@@ -106,6 +106,7 @@ class ReflectionContext
   typename super::StoredPointer target_future_adapter = 0;
   typename super::StoredPointer target_task_wait_throwing_resume_adapter = 0;
   typename super::StoredPointer target_task_future_wait_resume_adapter = 0;
+  bool supportsPriorityEscalation = false;
 
 public:
   using super::getBuilder;
@@ -164,7 +165,7 @@ class ReflectionContext
 
   ReflectionContext(const ReflectionContext &other) = delete;
   ReflectionContext &operator=(const ReflectionContext &other) = delete;
-  
+
   MemoryReader &getReader() {
     return *this->Reader;
   }
@@ -367,7 +368,7 @@ class ReflectionContext
 
         auto Begin = RemoteRef<void>(Addr, BufStart);
         auto Size = COFFSec->VirtualSize;
-        
+
         // FIXME: This code needs to be cleaned up and updated
         // to make it work for 32 bit platforms.
         Begin = Begin.atByteOffset(8);
@@ -605,7 +606,7 @@ class ReflectionContext
   }
 
   /// Parses metadata information from an ELF image. Because the Section
-  /// Header Table maybe be missing (for example, when reading from a 
+  /// Header Table maybe be missing (for example, when reading from a
   /// process) this method optionally receives a buffer with the contents
   /// of the image's file, from where it will the necessary information.
   ///
@@ -652,19 +653,19 @@ class ReflectionContext
     auto Magic = this->getReader().readBytes(ImageStart, sizeof(uint32_t));
     if (!Magic)
       return false;
-    
+
     uint32_t MagicWord;
     memcpy(&MagicWord, Magic.get(), sizeof(MagicWord));
-    
+
     // 32- and 64-bit Mach-O.
     if (MagicWord == llvm::MachO::MH_MAGIC) {
       return readMachOSections<MachOTraits<4>>(ImageStart);
     }
-    
+
     if (MagicWord == llvm::MachO::MH_MAGIC_64) {
       return readMachOSections<MachOTraits<8>>(ImageStart);
     }
-    
+
     // PE. (This just checks for the DOS header; `readPECOFF` will further
     // validate the existence of the PE header.)
     auto MagicBytes = (const char*)Magic.get();
@@ -739,7 +740,7 @@ class ReflectionContext
       return true;
     return ownsAddress(RemoteAddress(*MetadataAddress));
   }
-  
+
   /// Returns true if the address falls within a registered image.
   bool ownsAddressRaw(RemoteAddress Address) {
     for (auto Range : imageRanges) {
@@ -1193,7 +1194,7 @@ class ReflectionContext
     IterateConformanceTable(ConformancesAddr->getResolvedAddress(), Call);
     return llvm::None;
   }
-  
+
   /// Fetch the metadata pointer from a metadata allocation, or 0 if this
   /// allocation's tag is not handled or an error occurred.
   StoredPointer allocationMetadataPointer(
@@ -1291,7 +1292,7 @@ class ReflectionContext
       getReader().readPointer(AllocationPoolAddrAddr, sizeof(StoredPointer));
     if (!AllocationPoolAddr)
       return "failed to read value of " + AllocationPoolPointerName;
-    
+
     struct PoolRange {
       StoredPointer Begin;
       StoredSize Remaining;
@@ -1337,10 +1338,10 @@ class ReflectionContext
         Allocation.Ptr = RemoteAddr;
         Allocation.Size = Header->Size;
         Call(Allocation);
-        
+
         Offset += sizeof(AllocationHeader) + Header->Size;
       }
-      
+
       TrailerPtr = Trailer->PrevTrailer;
     }
     return llvm::None;
@@ -1421,13 +1422,20 @@ class ReflectionContext
 
   std::pair<llvm::Optional<std::string>, AsyncTaskInfo>
   asyncTaskInfo(StoredPointer AsyncTaskPtr) {
-    auto AsyncTaskObj = readObj<AsyncTask<Runtime>>(AsyncTaskPtr);
+    loadTargetPointers();
+
+    if (supportsPriorityEscalation) {
+      return {std::string("Failure reading async task with escalation support"), {}};
+    }
+
+    using AsyncTask = AsyncTask<Runtime, ActiveTaskStatusWithoutEscalation<Runtime>>;
+    auto AsyncTaskObj = readObj<AsyncTask>(AsyncTaskPtr);
     if (!AsyncTaskObj)
       return {std::string("failure reading async task"), {}};
 
     AsyncTaskInfo Info{};
     Info.JobFlags = AsyncTaskObj->Flags;
-    Info.TaskStatusFlags = AsyncTaskObj->PrivateStorage.Status.Flags;
+    Info.TaskStatusFlags = AsyncTaskObj->PrivateStorage.Status.Flags[0];
     Info.Id =
         AsyncTaskObj->Id | ((uint64_t)AsyncTaskObj->PrivateStorage.Id << 32);
     Info.AllocatorSlabPtr = AsyncTaskObj->PrivateStorage.Allocator.FirstSlab;
@@ -1461,7 +1469,7 @@ class ReflectionContext
         Info.ChildTasks.push_back(ChildTask);
 
         StoredPointer ChildFragmentAddr =
-            ChildTask + sizeof(AsyncTask<Runtime>);
+            ChildTask + sizeof(AsyncTask);
         auto ChildFragmentObj =
             readObj<ChildFragment<Runtime>>(ChildFragmentAddr);
         if (ChildFragmentObj)
@@ -1478,8 +1486,8 @@ class ReflectionContext
     // that's available.
     int IsCancelledFlag = 0x100;
     int IsRunningFlag = 0x800;
-    if (!(AsyncTaskObj->PrivateStorage.Status.Flags & IsCancelledFlag) &&
-        !(AsyncTaskObj->PrivateStorage.Status.Flags & IsRunningFlag)) {
+    if (!(AsyncTaskObj->PrivateStorage.Status.Flags[0] & IsCancelledFlag) &&
+        !(AsyncTaskObj->PrivateStorage.Status.Flags[0] & IsRunningFlag)) {
       auto ResumeContext = AsyncTaskObj->ResumeContextAndReserved[0];
       while (ResumeContext) {
         auto ResumeContextObj = readObj<AsyncContext<Runtime>>(ResumeContext);
@@ -1530,7 +1538,7 @@ class ReflectionContext
 private:
   // Get the most human meaningful "run job" function pointer from the task,
   // like AsyncTask::getResumeFunctionForLogging does.
-  StoredPointer getRunJob(const AsyncTask<Runtime> *AsyncTaskObj) {
+  StoredPointer getRunJob(const AsyncTask<Runtime, ActiveTaskStatusWithoutEscalation<Runtime>> *AsyncTaskObj) {
     auto Fptr = stripSignedPointer(AsyncTaskObj->RunJob);
 
     loadTargetPointers();
@@ -1589,6 +1597,11 @@ class ReflectionContext
         getFunc("_swift_concurrency_debug_task_wait_throwing_resume_adapter");
     target_task_future_wait_resume_adapter =
         getFunc("_swift_concurrency_debug_task_future_wait_resume_adapter");
+    auto supportsPriorityEscalationAddr = getReader().getSymbolAddress("_swift_concurrency_debug_supportsPriorityEscalation");
+    if (supportsPriorityEscalationAddr) {
+      getReader().readInteger(supportsPriorityEscalationAddr, &supportsPriorityEscalation);
+    }
+
     setupTargetPointers = true;
   }
 

include/swift/Reflection/RuntimeInternals.h

@@ -94,28 +94,34 @@ struct StackAllocator {
 };
 
 template <typename Runtime>
-struct ActiveTaskStatus {
+struct ActiveTaskStatusWithEscalation {
+  uint32_t Flags;
+  uint32_t DrainLock[(sizeof(typename Runtime::StoredPointer) == 8) ? 1 : 2];
   typename Runtime::StoredPointer Record;
-  typename Runtime::StoredSize Flags;
 };
 
 template <typename Runtime>
+struct ActiveTaskStatusWithoutEscalation {
+  uint32_t Flags[sizeof(typename Runtime::StoredPointer) == 8 ? 2 : 1];
+  typename Runtime::StoredPointer Record;
+};
+
+template <typename Runtime, typename ActiveTaskStatus>
 struct AsyncTaskPrivateStorage {
-  ActiveTaskStatus<Runtime> Status;
+  ActiveTaskStatus Status;
   StackAllocator<Runtime> Allocator;
   typename Runtime::StoredPointer Local;
   typename Runtime::StoredPointer ExclusivityAccessSet[2];
   uint32_t Id;
 };
 
-template <typename Runtime>
+template <typename Runtime, typename ActiveTaskStatus>
 struct AsyncTask: Job<Runtime> {
-  // On 64-bit, there's a Reserved64 after ResumeContext.  
+  // On 64-bit, there's a Reserved64 after ResumeContext.
   typename Runtime::StoredPointer ResumeContextAndReserved[
     sizeof(typename Runtime::StoredPointer) == 8 ? 2 : 1];
-
   union {
-    AsyncTaskPrivateStorage<Runtime> PrivateStorage;
+    AsyncTaskPrivateStorage<Runtime, ActiveTaskStatus> PrivateStorage;
     typename Runtime::StoredPointer PrivateStorageRaw[14];
   };
 };

include/swift/Runtime/Concurrency.h

@@ -41,6 +41,17 @@
 #endif
 #endif
 
+// Does the runtime provide priority escalation support?
+#ifndef SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
+#if SWIFT_CONCURRENCY_ENABLE_DISPATCH && \
+    __has_include(<dispatch/swift_concurrency_private.h>) && __APPLE__ && \
+    SWIFT_POINTER_IS_8_BYTES
+#define SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION 1
+#else
+#define SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION 0
+#endif
+#endif /* SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION */
+
 namespace swift {
 class DefaultActor;
 class TaskOptionRecord;
@@ -616,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/Debug.h

@@ -45,6 +45,10 @@ const void *const _swift_concurrency_debug_task_wait_throwing_resume_adapter;
 SWIFT_EXPORT_FROM(swift_Concurrency)
 const void *const _swift_concurrency_debug_task_future_wait_resume_adapter;
 
+/// Whether the runtime we are inspecting supports priority escalation
+SWIFT_EXPORT_FROM(swift_Concurrency)
+bool _swift_concurrency_debug_supportsPriorityEscalation;
+
 } // namespace swift
 
 #endif