Task Escalation support on arm64_32 Apple platforms

Radar-Id: rdar://problem/88600541

Author: rokhinip

include/swift/Reflection/RuntimeInternals.h

@@ -108,11 +108,12 @@ struct ActiveTaskStatusWithoutEscalation {
 
 template <typename Runtime, typename ActiveTaskStatus>
 struct AsyncTaskPrivateStorage {
+  typename Runtime::StoredPointer ExclusivityAccessSet[2];
   ActiveTaskStatus Status;
   StackAllocator<Runtime> Allocator;
   typename Runtime::StoredPointer Local;
-  typename Runtime::StoredPointer ExclusivityAccessSet[2];
   uint32_t Id;
+  uint32_t BasePriority;
 };
 
 template <typename Runtime, typename ActiveTaskStatus>

include/swift/Runtime/Concurrency.h

@@ -45,7 +45,7 @@
 #ifndef SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
 #if SWIFT_CONCURRENCY_ENABLE_DISPATCH && \
     __has_include(<dispatch/swift_concurrency_private.h>) && __APPLE__ && \
-    SWIFT_POINTER_IS_8_BYTES
+    (defined(__arm64__) || defined(__x86_64__))
 #define SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION 1
 #else
 #define SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION 0

stdlib/public/Concurrency/Task.cpp

@@ -148,7 +148,7 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
 
       // Escalate the priority of this task based on the priority
       // of the waiting task.
-      auto status = waitingTask->_private().Status.load(std::memory_order_relaxed);
+      auto status = waitingTask->_private()._status().load(std::memory_order_relaxed);
       swift_task_escalate(this, status.getStoredPriority());
 
       _swift_task_clearCurrent();
@@ -491,7 +491,7 @@ const void *AsyncTask::getResumeFunctionForLogging() {
 JobPriority swift::swift_task_currentPriority(AsyncTask *task)
 {
   // This is racey but this is to be used in an API is inherently racey anyways.
-  auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
   return oldStatus.getStoredPriority();
 }
 

stdlib/public/Concurrency/TaskPrivate.h

@@ -229,7 +229,28 @@ class TaskFutureWaitAsyncContext : public AsyncContext {
 /// only need to do double wide atomics if we need to reach for the
 /// StatusRecord pointers and therefore have to update the flags at the same
 /// time.
-class alignas(sizeof(void*) * 2) ActiveTaskStatus {
+///
+/// Size requirements:
+///     On 64 bit systems or if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION=1,
+///     the field is 16 bytes long.
+///
+///     Otherwise, it is 8 bytes long.
+///
+/// Alignment requirements:
+///     On 64 bit systems or if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION=1,
+///     this 16-byte field needs to be 16 byte aligned to be able to do aligned
+///     atomic stores field.
+///
+///     On all other systems, it needs to be 8 byte aligned for the atomic
+///     stores.
+///
+///     As a result of varying alignment needs, we've marked the class as
+///     needing 2-word alignment but on arm64_32 with
+///     SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION=1, 16 byte alignment is
+///     achieved through careful arrangement of the storage for this in the
+///     AsyncTask::PrivateStorage. The additional alignment requirements are
+///     enforced by static asserts below.
+class alignas(2 * sizeof(void*)) ActiveTaskStatus {
   enum : uint32_t {
     /// The max priority of the task. This is always >= basePriority in the task
     PriorityMask = 0xFF,
@@ -463,12 +484,12 @@ class alignas(sizeof(void*) * 2) ActiveTaskStatus {
 };
 
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION && SWIFT_POINTER_IS_4_BYTES
-static_assert(sizeof(ActiveTaskStatus) == 4 * sizeof(uintptr_t),
-  "ActiveTaskStatus is 4 words large");
+#define ACTIVE_TASK_STATUS_SIZE (4 * (sizeof(uintptr_t)))
 #else
-static_assert(sizeof(ActiveTaskStatus) == 2 * sizeof(uintptr_t),
-  "ActiveTaskStatus is 2 words large");
+#define ACTIVE_TASK_STATUS_SIZE (2 * (sizeof(uintptr_t)))
 #endif
+static_assert(sizeof(ActiveTaskStatus) == ACTIVE_TASK_STATUS_SIZE,
+  "ActiveTaskStatus is of incorrect size");
 
 /// The size of an allocator slab. We want the full allocation to fit into a
 /// 1024-byte malloc quantum. We subtract off the slab header size, plus a
@@ -481,9 +502,16 @@ using TaskAllocator = StackAllocator<SlabCapacity, &TaskAllocatorSlabMetadata>;
 
 /// Private storage in an AsyncTask object.
 struct AsyncTask::PrivateStorage {
-  /// The currently-active information about cancellation.
-  /// Currently two words.
-  swift::atomic<ActiveTaskStatus> Status;
+  /// State inside the AsyncTask whose state is only managed by the exclusivity
+  /// runtime in stdlibCore. We zero initialize to provide a safe initial value,
+  /// but actually initialize its bit state to a const global provided by
+  /// libswiftCore so that libswiftCore can control the layout of our initial
+  /// state.
+  uintptr_t ExclusivityAccessSet[2] = {0, 0};
+
+  /// Storage for the ActiveTaskStatus. See doc for ActiveTaskStatus for size
+  /// and alignment requirements.
+  alignas(2 * sizeof(void *)) char StatusStorage[sizeof(ActiveTaskStatus)];
 
   /// The allocator for the task stack.
   /// Currently 2 words + 8 bytes.
@@ -493,13 +521,6 @@ struct AsyncTask::PrivateStorage {
   /// Currently one word.
   TaskLocal::Storage Local;
 
-  /// State inside the AsyncTask whose state is only managed by the exclusivity
-  /// runtime in stdlibCore. We zero initialize to provide a safe initial value,
-  /// but actually initialize its bit state to a const global provided by
-  /// libswiftCore so that libswiftCore can control the layout of our initial
-  /// state.
-  uintptr_t ExclusivityAccessSet[2] = {0, 0};
-
   /// The top 32 bits of the task ID. The bottom 32 bits are in Job::Id.
   uint32_t Id;
 
@@ -514,19 +535,22 @@ struct AsyncTask::PrivateStorage {
   // Always create an async task with max priority in ActiveTaskStatus = base
   // priority. It will be updated later if needed.
   PrivateStorage(JobPriority basePri)
-      : Status(ActiveTaskStatus(basePri)), Local(TaskLocal::Storage()),
-        BasePriority(basePri) {}
+      : Local(TaskLocal::Storage()), BasePriority(basePri) {
+    _status().store(ActiveTaskStatus(basePri), std::memory_order_relaxed);
+  }
 
   PrivateStorage(JobPriority basePri, void *slab, size_t slabCapacity)
-      : Status(ActiveTaskStatus(basePri)), Allocator(slab, slabCapacity),
-        Local(TaskLocal::Storage()), BasePriority(basePri) {}
+      : Allocator(slab, slabCapacity), Local(TaskLocal::Storage()),
+          BasePriority(basePri) {
+    _status().store(ActiveTaskStatus(basePri), std::memory_order_relaxed);
+  }
 
   /// Called on the thread that was previously executing the task that we are
   /// now trying to complete.
   void complete(AsyncTask *task) {
     // Drain unlock the task and remove any overrides on thread as a
     // result of the task
-    auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
+    auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
     while (true) {
       // Task is completing, it shouldn't have any records and therefore
       // cannot be status record locked.
@@ -542,7 +566,7 @@ struct AsyncTask::PrivateStorage {
 
       // This can fail since the task can still get concurrently cancelled or
       // escalated.
-      if (task->_private().Status.compare_exchange_weak(oldStatus, newStatus,
+      if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
               /* success */ std::memory_order_relaxed,
               /* failure */ std::memory_order_relaxed)) {
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
@@ -561,12 +585,21 @@ struct AsyncTask::PrivateStorage {
 
     this->~PrivateStorage();
   }
+
+  swift::atomic<ActiveTaskStatus> &_status() {
+    return reinterpret_cast<swift::atomic<ActiveTaskStatus>&> (this->StatusStorage);
+  }
+
+  const swift::atomic<ActiveTaskStatus> &_status() const {
+    return reinterpret_cast<const swift::atomic<ActiveTaskStatus>&> (this->StatusStorage);
+  }
 };
 
-static_assert(sizeof(AsyncTask::PrivateStorage)
-                <= sizeof(AsyncTask::OpaquePrivateStorage) &&
-              alignof(AsyncTask::PrivateStorage)
-                <= alignof(AsyncTask::OpaquePrivateStorage),
+// It will be aligned to 2 words on all platforms. On arm64_32, we have an
+// additional requirement where it is aligned to 4 words.
+static_assert(((offsetof(AsyncTask, Private) + offsetof(AsyncTask::PrivateStorage, StatusStorage)) % ACTIVE_TASK_STATUS_SIZE == 0),
+   "StatusStorage is not aligned in the AsyncTask");
+static_assert(sizeof(AsyncTask::PrivateStorage) <= sizeof(AsyncTask::OpaquePrivateStorage),
               "Task-private storage doesn't fit in reserved space");
 
 inline AsyncTask::PrivateStorage &
@@ -599,7 +632,7 @@ inline const AsyncTask::PrivateStorage &AsyncTask::_private() const {
 }
 
 inline bool AsyncTask::isCancelled() const {
-  return _private().Status.load(std::memory_order_relaxed)
+  return _private()._status().load(std::memory_order_relaxed)
                           .isCancelled();
 }
 
@@ -612,7 +645,7 @@ inline void AsyncTask::flagAsRunning() {
   qos_class_t overrideFloor = threadOverrideInfo.override_qos_floor;
 retry:;
 #endif
-  auto oldStatus = _private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = _private()._status().load(std::memory_order_relaxed);
   while (true) {
     // We can get here from being suspended or being enqueued
     assert(!oldStatus.isRunning());
@@ -636,7 +669,7 @@ retry:;
     newStatus = newStatus.withoutStoredPriorityEscalation();
     newStatus = newStatus.withoutEnqueued();
 
-    if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
+    if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
              /* success */ std::memory_order_relaxed,
              /* failure */ std::memory_order_relaxed)) {
       newStatus.traceStatusChanged(this);
@@ -667,7 +700,7 @@ retry:;
 inline void AsyncTask::flagAsEnqueuedOnExecutor(ExecutorRef newExecutor) {
 
   SWIFT_TASK_DEBUG_LOG("%p->flagAsEnqueuedOnExecutor()", this);
-  auto oldStatus = _private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = _private()._status().load(std::memory_order_relaxed);
   auto newStatus = oldStatus;
 
   while (true) {
@@ -682,7 +715,7 @@ inline void AsyncTask::flagAsEnqueuedOnExecutor(ExecutorRef newExecutor) {
     newStatus = newStatus.withoutStoredPriorityEscalation();
     newStatus = newStatus.withEnqueued();
 
-    if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
+    if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
             /* success */std::memory_order_relaxed,
             /* failure */std::memory_order_relaxed)) {
       break;
@@ -713,7 +746,7 @@ inline void AsyncTask::flagAsEnqueuedOnExecutor(ExecutorRef newExecutor) {
 
 inline void AsyncTask::flagAsSuspended() {
   SWIFT_TASK_DEBUG_LOG("%p->flagAsSuspended()", this);
-  auto oldStatus = _private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = _private()._status().load(std::memory_order_relaxed);
   auto newStatus = oldStatus;
   while (true) {
     // We can only be suspended if we were previously running. See state
@@ -723,7 +756,7 @@ inline void AsyncTask::flagAsSuspended() {
     newStatus = oldStatus.withRunning(false);
     newStatus = newStatus.withoutStoredPriorityEscalation();
 
-    if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
+    if (_private()._status().compare_exchange_weak(oldStatus, newStatus,
             /* success */std::memory_order_relaxed,
             /* failure */std::memory_order_relaxed)) {
       break;

stdlib/public/Concurrency/TaskStatus.cpp

@@ -96,7 +96,7 @@ static void waitForStatusRecordUnlock(AsyncTask *task,
 
     bool waited = waiter.tryReloadAndWait([&]() -> StatusRecordLockRecord* {
       // Check that oldStatus is still correct.
-      oldStatus = task->_private().Status.load(std::memory_order_acquire);
+      oldStatus = task->_private()._status().load(std::memory_order_acquire);
       if (!oldStatus.isStatusRecordLocked())
         return nullptr;
 
@@ -109,7 +109,7 @@ static void waitForStatusRecordUnlock(AsyncTask *task,
       return;
 
     // Reload the status before trying to relock.
-    oldStatus = task->_private().Status.load(std::memory_order_acquire);
+    oldStatus = task->_private()._status().load(std::memory_order_acquire);
     if (!oldStatus.isStatusRecordLocked())
       return;
   }
@@ -162,7 +162,7 @@ static bool withStatusRecordLock(AsyncTask *task,
 
     // Install the lock record as the top of the queue.
     ActiveTaskStatus newStatus = status.withLockingRecord(lockingRecord);
-    if (task->_private().Status.compare_exchange_weak(status, newStatus,
+    if (task->_private()._status().compare_exchange_weak(status, newStatus,
             /*success*/ std::memory_order_release,
             /*failure*/ loadOrdering)) {
       status = newStatus;
@@ -187,7 +187,7 @@ static bool withStatusRecordLock(AsyncTask *task,
     assert(status.isStatusRecordLocked());
     auto newStatus = status.withoutLockingRecord();
 
-    if (task->_private().Status.compare_exchange_weak(status, newStatus,
+    if (task->_private()._status().compare_exchange_weak(status, newStatus,
             /*success*/ std::memory_order_relaxed,
             /*failure*/ std::memory_order_relaxed)) {
       status.traceStatusChanged(task);
@@ -208,7 +208,7 @@ static bool withStatusRecordLock(AsyncTask *task,
                                  LockContext lockContext,
                                  Fn &&fn) {
   ActiveTaskStatus status =
-    task->_private().Status.load(getLoadOrdering(lockContext));
+    task->_private()._status().load(getLoadOrdering(lockContext));
   return withStatusRecordLock(task, lockContext, status, [&] {
     fn(status);
   });
@@ -226,7 +226,7 @@ bool swift::addStatusRecord(
   auto task = swift_task_getCurrent();
   // Load the current state. We can use a relaxed load because we're
   // synchronous with the task.
-  auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
 
   while (true) {
     // Wait for any active lock to be released.
@@ -243,7 +243,7 @@ bool swift::addStatusRecord(
       // We have to use a release on success to make the initialization of
       // the new record visible to an asynchronous thread trying to modify the
       // status records
-      if (task->_private().Status.compare_exchange_weak(oldStatus, newStatus,
+      if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
               /*success*/ std::memory_order_release,
               /*failure*/ std::memory_order_relaxed)) {
         return true;
@@ -263,7 +263,7 @@ bool swift::removeStatusRecord(TaskStatusRecord *record) {
                        record, task);
 
   // Load the current state.
-  auto &status = task->_private().Status;
+  auto &status = task->_private()._status();
   auto oldStatus = status.load(std::memory_order_relaxed);
 
   while (true) {
@@ -358,7 +358,7 @@ void swift::updateNewChildWithParentAndGroupState(AsyncTask *child,
   // cannot be accessed by anyone else since it hasn't been linked in yet.
   // Avoids the extra logic in `swift_task_cancel` and `swift_task_escalate`
   auto oldChildTaskStatus =
-      child->_private().Status.load(std::memory_order_relaxed);
+      child->_private()._status().load(std::memory_order_relaxed);
   assert(oldChildTaskStatus.getInnermostRecord() == NULL);
 
   auto newChildTaskStatus = oldChildTaskStatus;
@@ -372,7 +372,7 @@ void swift::updateNewChildWithParentAndGroupState(AsyncTask *child,
   newChildTaskStatus =
       newChildTaskStatus.withNewPriority(withUserInteractivePriorityDowngrade(pri));
 
-  child->_private().Status.store(newChildTaskStatus, std::memory_order_relaxed);
+  child->_private()._status().store(newChildTaskStatus, std::memory_order_relaxed);
 }
 
 SWIFT_CC(swift)
@@ -451,7 +451,7 @@ SWIFT_CC(swift)
 static void swift_task_cancelImpl(AsyncTask *task) {
   SWIFT_TASK_DEBUG_LOG("cancel task = %p", task);
 
-  auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
   auto newStatus = oldStatus;
   while (true) {
     if (oldStatus.isCancelled()) {
@@ -461,7 +461,7 @@ static void swift_task_cancelImpl(AsyncTask *task) {
     // Set cancelled bit even if oldStatus.isStatusRecordLocked()
     newStatus = oldStatus.withCancelled();
 
-    if (task->_private().Status.compare_exchange_weak(oldStatus, newStatus,
+    if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
             /*success*/ std::memory_order_relaxed,
             /*failure*/ std::memory_order_relaxed)) {
       break;
@@ -550,7 +550,7 @@ JobPriority
 static swift_task_escalateImpl(AsyncTask *task, JobPriority newPriority) {
 
   SWIFT_TASK_DEBUG_LOG("Escalating %p to %#x priority", task, newPriority);
-  auto oldStatus = task->_private().Status.load(std::memory_order_relaxed);
+  auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
   auto newStatus = oldStatus;
 
   while (true) {
@@ -569,7 +569,7 @@ static swift_task_escalateImpl(AsyncTask *task, JobPriority newPriority) {
       newStatus = oldStatus.withNewPriority(newPriority);
     }
 
-    if (task->_private().Status.compare_exchange_weak(oldStatus, newStatus,
+    if (task->_private()._status().compare_exchange_weak(oldStatus, newStatus,
             /* success */ std::memory_order_relaxed,
             /* failure */ std::memory_order_relaxed)) {
       break;
@@ -582,7 +582,7 @@ static swift_task_escalateImpl(AsyncTask *task, JobPriority newPriority) {
     ActiveTaskStatus *taskStatus;
     dispatch_lock_t *executionLock;
 
-    taskStatus = (ActiveTaskStatus *) &task->_private().Status;
+    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);
@@ -631,7 +631,7 @@ static NearestTaskDeadline swift_task_getNearestDeadlineImpl(AsyncTask *task) {
   // ignoring the possibility of a concurrent cancelling task.
 
   // Load the current state.
-  auto &status = task->_private().Status;
+  auto &status = task->_private()._status();
   auto oldStatus = status.load(std::memory_order_relaxed);
 
   NearestTaskDeadline result;