Provide a convenience for addStatusRecord for clients who have already

done the load or who need the oldStatus information after adding the
status record.

Change some of the memory barrier logic since we can take advantage of
load-through HW address dependency.

Radar-Id: rdar://problem/105634683

Author: rokhinip

include/swift/ABI/TaskStatus.h

@@ -32,16 +32,14 @@ namespace swift {
 /// context), partially initialized, and then atomically added to the task with
 /// `addStatusRecord`.
 ///
-/// Status records currently are only added to a task by itself but they may be
-/// removed from the task by other threads, or accessed by other threads
-/// asynchronous to the task for other purposes.
-///
-/// As a result, while registered with the task, a status record should only be
-/// modified in ways that respect the possibility of asynchronous access by a
-/// different thread.  In particular, the chain of status records must not be
-/// disturbed. When the task leaves the scope that requires the status record,
-/// the record can be unregistered from the task with `removeStatusRecord`, at
-/// which point the memory can be returned to the system.
+/// Status records can be added to or removed from  a task by itself or by other
+/// threads. As a result, while registered with the task, a status record
+/// should only be modified in ways that respect the possibility of asynchronous
+/// access by a different thread.  In particular, the chain of status records
+/// must not be disturbed. When the task leaves the scope that requires the
+/// status record, the record can be unregistered from the task with
+/// `removeStatusRecord`, at which point the memory can be returned to the
+/// system.
 class TaskStatusRecord {
 public:
   TaskStatusRecordFlags Flags;

include/swift/Runtime/Atomic.h

@@ -25,6 +25,7 @@
 #include <intrin.h>
 #endif
 
+
 // FIXME: Workaround for rdar://problem/18889711. 'Consume' does not require
 // a barrier on ARM64, but LLVM doesn't know that. Although 'relaxed'
 // is formally UB by C++11 language rules, we should be OK because neither
@@ -36,6 +37,7 @@
 #  define SWIFT_MEMORY_ORDER_CONSUME (std::memory_order_consume)
 #endif
 
+
 #if defined(_M_ARM) || defined(__arm__) || defined(__aarch64__)
 #define SWIFT_HAS_MSVC_ARM_ATOMICS 1
 #else

stdlib/public/Concurrency/AsyncLet.cpp

@@ -48,16 +48,16 @@ class alignas(Alignment_AsyncLet) AsyncLetImpl: public ChildTaskStatusRecord {
     HasResult = 1 << 0,
     DidAllocateFromParentTask = 1 << 1,
   };
-  
+
   /// The task that was kicked off to initialize this `async let`,
   /// and flags.
   llvm::PointerIntPair<AsyncTask *, 2, unsigned> taskAndFlags;
-  
+
   /// Reserved space for a future_wait context frame, used during suspensions
   /// on the child task future.
   std::aligned_storage<sizeof(TaskFutureWaitAsyncContext),
                        alignof(TaskFutureWaitAsyncContext)>::type futureWaitContextStorage;
-  
+
   friend class ::swift::AsyncTask;
 
 public:
@@ -77,43 +77,43 @@ class alignas(Alignment_AsyncLet) AsyncLetImpl: public ChildTaskStatusRecord {
   AsyncTask *getTask() const {
     return taskAndFlags.getPointer();
   }
-  
+
   bool hasResultInBuffer() const {
     return taskAndFlags.getInt() & HasResult;
   }
-  
+
   void setHasResultInBuffer(bool value = true) {
     if (value)
       taskAndFlags.setInt(taskAndFlags.getInt() | HasResult);
     else
       taskAndFlags.setInt(taskAndFlags.getInt() & ~HasResult);
   }
-  
+
   bool didAllocateFromParentTask() const {
     return taskAndFlags.getInt() & DidAllocateFromParentTask;
   }
-  
+
   void setDidAllocateFromParentTask(bool value = true) {
     if (value)
       taskAndFlags.setInt(taskAndFlags.getInt() | DidAllocateFromParentTask);
     else
       taskAndFlags.setInt(taskAndFlags.getInt() & ~DidAllocateFromParentTask);
   }
-  
+
   // The compiler preallocates a large fixed space for the `async let`, with the
   // intent that most of it be used for the child task context. The next two
   // methods return the address and size of that space.
-  
+
   /// Return a pointer to the unused space within the async let block.
   void *getPreallocatedSpace() {
     return (void*)(this + 1);
   }
-  
+
   /// Return the size of the unused space within the async let block.
   static constexpr size_t getSizeOfPreallocatedSpace() {
     return sizeof(AsyncLet) - sizeof(AsyncLetImpl);
   }
-  
+
   TaskFutureWaitAsyncContext *getFutureContext() {
     return reinterpret_cast<TaskFutureWaitAsyncContext*>(&futureWaitContextStorage);
   }
@@ -149,11 +149,11 @@ void swift::asyncLet_addImpl(AsyncTask *task, AsyncLet *asyncLet,
 
   // ok, now that the async let task actually is initialized: attach it to the
   // current task
-  bool addedRecord =
-      addStatusRecord(record, [&](ActiveTaskStatus parentStatus, ActiveTaskStatus& newStatus) {
-        updateNewChildWithParentAndGroupState(task, parentStatus, NULL);
-        return true;
-      });
+  bool addedRecord = addStatusRecordToSelf(record,
+      [&](ActiveTaskStatus parentStatus, ActiveTaskStatus& newStatus) {
+    updateNewChildWithParentAndGroupState(task, parentStatus, NULL);
+    return true;
+  });
   (void)addedRecord;
   assert(addedRecord);
 }
@@ -279,13 +279,13 @@ static void _asyncLet_get_throwing_continuation(
         SWIFT_CONTEXT void *error) {
   auto continuationContext = static_cast<AsyncLetContinuationContext*>(callContext);
   auto alet = continuationContext->alet;
-  
+
   // If the future completed successfully, its result is now in the async let
   // buffer.
   if (!error) {
     asImpl(alet)->setHasResultInBuffer();
   }
-  
+
   // Continue the caller's execution.
   auto throwingResume
     = reinterpret_cast<ThrowingTaskFutureWaitContinuationFunction*>(callContext->ResumeParent);
@@ -303,14 +303,14 @@ static void swift_asyncLet_get_throwingImpl(
   if (asImpl(alet)->hasResultInBuffer()) {
     return resumeFunction(callerContext, nullptr);
   }
-  
+
   auto aletContext = static_cast<AsyncLetContinuationContext*>(callContext);
   aletContext->ResumeParent
     = reinterpret_cast<TaskContinuationFunction*>(resumeFunction);
   aletContext->Parent = callerContext;
   aletContext->alet = alet;
   auto futureContext = asImpl(alet)->getFutureContext();
-  
+
   // Unlike the non-throwing variant, whether we end up with a result depends
   // on the success of the task. If we raise an error, then the result buffer
   // will not be populated. Save the async let binding so we can fetch it
@@ -375,7 +375,7 @@ static void asyncLet_finish_after_task_completion(SWIFT_ASYNC_CONTEXT AsyncConte
   if (alet->didAllocateFromParentTask()) {
     swift_task_dealloc(task);
   }
-  
+
   return reinterpret_cast<ThrowingTaskFutureWaitContinuationFunction*>(resumeFunction)
     (callerContext, error);
 }
@@ -390,7 +390,7 @@ static void _asyncLet_finish_continuation(
     = reinterpret_cast<AsyncLetContinuationContext*>(callContext);
   auto alet = continuationContext->alet;
   auto resultBuffer = continuationContext->resultBuffer;
-  
+
   // Destroy the error, or the result that was stored to the buffer.
   if (error) {
     swift_errorRelease((SwiftError*)error);
@@ -436,7 +436,7 @@ static void swift_asyncLet_finishImpl(SWIFT_ASYNC_CONTEXT AsyncContext *callerCo
   aletContext->alet = alet;
   aletContext->resultBuffer = reinterpret_cast<OpaqueValue*>(resultBuffer);
   auto futureContext = asImpl(alet)->getFutureContext();
-  
+
   // TODO: It would be nice if we could await the future without having to
   // provide a buffer to store the value to, since we're going to dispose of
   // it anyway.
@@ -504,7 +504,7 @@ static void _asyncLet_consume_throwing_continuation(
   // Get the async let pointer so we can destroy the task.
   auto continuationContext = static_cast<AsyncLetContinuationContext*>(callContext);
   auto alet = continuationContext->alet;
-  
+
   return asyncLet_finish_after_task_completion(callContext->Parent,
                                                alet,
                                                callContext->ResumeParent,
@@ -528,14 +528,14 @@ static void swift_asyncLet_consume_throwingImpl(
                    callContext,
                    nullptr);
   }
-  
+
   auto aletContext = static_cast<AsyncLetContinuationContext*>(callContext);
   aletContext->ResumeParent
     = reinterpret_cast<TaskContinuationFunction*>(resumeFunction);
   aletContext->Parent = callerContext;
   aletContext->alet = alet;
   auto futureContext = asImpl(alet)->getFutureContext();
-  
+
   // Unlike the non-throwing variant, whether we end up with a result depends
   // on the success of the task. If we raise an error, then the result buffer
   // will not be populated. Save the async let binding so we can fetch it

stdlib/public/Concurrency/Task.cpp

@@ -117,7 +117,6 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
   // NOTE: this acquire synchronizes with `completeFuture`.
   auto queueHead = fragment->waitQueue.load(std::memory_order_acquire);
   bool contextInitialized = false;
-  auto escalatedPriority = JobPriority::Unspecified;
   while (true) {
     switch (queueHead.getStatus()) {
     case Status::Error:
@@ -192,7 +191,6 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
 void NullaryContinuationJob::process(Job *_job) {
   auto *job = cast<NullaryContinuationJob>(_job);
 
-  auto *task = job->Task;
   auto *continuation = job->Continuation;
 
   delete job;
@@ -1496,7 +1494,7 @@ swift_task_addCancellationHandlerImpl(
 
   bool fireHandlerNow = false;
 
-  addStatusRecord(record, [&](ActiveTaskStatus oldStatus, ActiveTaskStatus& newStatus) {
+  addStatusRecordToSelf(record, [&](ActiveTaskStatus oldStatus, ActiveTaskStatus& newStatus) {
     if (oldStatus.isCancelled()) {
       fireHandlerNow = true;
       // We don't fire the cancellation handler here since this function needs

stdlib/public/Concurrency/TaskGroup.cpp

@@ -921,7 +921,7 @@ static void swift_taskGroup_initializeWithFlagsImpl(size_t rawGroupFlags,
   assert(record->getKind() == swift::TaskStatusRecordKind::TaskGroup);
 
   // ok, now that the group actually is initialized: attach it to the task
-  addStatusRecord(record, [&](ActiveTaskStatus oldStatus, ActiveTaskStatus& newStatus) {
+  addStatusRecordToSelf(record, [&](ActiveTaskStatus oldStatus, ActiveTaskStatus& newStatus) {
     // If the task has already been cancelled, reflect that immediately in
     // the group's status.
     if (oldStatus.isCancelled()) {

stdlib/public/Concurrency/TaskPrivate.h

@@ -147,8 +147,13 @@ void removeStatusRecord(AsyncTask *task, TaskStatusRecord *record,
 SWIFT_CC(swift)
 void removeStatusRecord(TaskStatusRecord *record);
 
-/// Add a status record to the current task. This must be called synchronously
-/// with the task.
+/// Add a status record to the input task.
+///
+/// Clients can optionally pass in the status of the task if they have already
+/// done the load on it already or if they require the oldStatus on the task
+/// prior to the atomic ActiveTaskStatus update that addStatusRecord will
+/// perform. This status will be updated with the last status on the task prior
+/// to updating it with the new status if the input function_ref allows so.
 ///
 /// This function also takes in a function_ref which is given the old
 /// ActiveTaskStatus on the task and a reference to the new ActiveTaskStatus
@@ -162,9 +167,27 @@ void removeStatusRecord(TaskStatusRecord *record);
 /// If the function_ref returns false, the status record is not added to the
 /// task. This function_ref may be called multiple times and must be idempotent.
 /// The new status passed to `fn` is freshly derived from the current status and
-/// does not include modifications made by previous runs through the loop
+/// does not include modifications made by previous runs through the loop.
+///
+/// A convenience overload is also provided for clients who have not done the
+/// load on the task status prior to adding the record
 SWIFT_CC(swift)
-bool addStatusRecord(TaskStatusRecord *record,
+bool addStatusRecord(AsyncTask *task, TaskStatusRecord *record,
+     llvm::function_ref<bool(ActiveTaskStatus, ActiveTaskStatus&)> testAddRecord);
+
+SWIFT_CC(swift)
+bool addStatusRecord(AsyncTask *task, TaskStatusRecord *record,
+     ActiveTaskStatus& taskStatus,
+     llvm::function_ref<bool(ActiveTaskStatus, ActiveTaskStatus&)> testAddRecord);
+
+/// Convenience functions for clients who want to just add a record to the task
+/// on the current thread
+SWIFT_CC(swift)
+bool addStatusRecordToSelf(TaskStatusRecord *record,
+     llvm::function_ref<bool(ActiveTaskStatus, ActiveTaskStatus&)> testAddRecord);
+
+SWIFT_CC(swift)
+bool addStatusRecordToSelf(TaskStatusRecord *record,  ActiveTaskStatus& taskStatus,
      llvm::function_ref<bool(ActiveTaskStatus, ActiveTaskStatus&)> testAddRecord);
 
 /// A helper function for updating a new child task that is created with
@@ -958,17 +981,13 @@ void AsyncTask::flagAsSuspended(TaskDependencyStatusRecord *dependencyStatusReco
                   dependencyStatusRecord);
   assert(dependencyStatusRecord != NULL);
 
-  bool taskWasEscalated = false;
-  JobPriority taskStoredPriority = JobPriority::Unspecified;
+  auto oldStatus = _private()._status().load(std::memory_order_relaxed);
+  // We can only be suspended if we were previously running. See state
+  // transitions listed out in Task.h
+  assert(oldStatus.isRunning() && !oldStatus.isEnqueued());
 
-  addStatusRecord(dependencyStatusRecord, [&](ActiveTaskStatus oldStatus,
+  addStatusRecord(this, dependencyStatusRecord, oldStatus, [&](ActiveTaskStatus unused,
                   ActiveTaskStatus &newStatus) {
-    // We can only be suspended if we were previously running. See state
-    // transitions listed out in Task.h
-    assert(oldStatus.isRunning() && !oldStatus.isEnqueued());
-    taskWasEscalated = oldStatus.isStoredPriorityEscalated();
-    taskStoredPriority = oldStatus.getStoredPriority();
-
     newStatus = newStatus.withRunning(false);
     newStatus = newStatus.withoutStoredPriorityEscalation();
     newStatus = newStatus.withTaskDependency();
@@ -985,7 +1004,7 @@ void AsyncTask::flagAsSuspended(TaskDependencyStatusRecord *dependencyStatusReco
     // published it in the ActiveTaskStatus.
     SWIFT_TASK_DEBUG_LOG("[Dependency] Escalate the dependency %p of task %p",
                   dependencyStatusRecord, this);
-    dependencyStatusRecord->performEscalationAction(taskStoredPriority);
+    dependencyStatusRecord->performEscalationAction(newStatus.getStoredPriority());
 
     // Always add the dependency status record
     return true;
@@ -994,10 +1013,10 @@ void AsyncTask::flagAsSuspended(TaskDependencyStatusRecord *dependencyStatusReco
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
   // Successfully dropped task drain lock, make sure to remove override on
   // thread due to task
-  if (taskWasEscalated) {
+  if (oldStatus.isStoredPriorityEscalated()) {
      SWIFT_TASK_DEBUG_LOG("[Override] Reset override %#x on thread from task %p",
-       taskStoredPriority, this);
-     swift_dispatch_lock_override_end((qos_class_t) taskStoredPriority);
+       oldStatus.getStoredPriority(), this);
+     swift_dispatch_lock_override_end((qos_class_t) oldStatus.getStoredPriority());
   }
 #endif