Merge pull request swiftlang#63800 from apple/rokhinip/105634683-allow-addStatusRecord-to-other-task

Loosen requirement that task can only addStatusRecords to itself

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