back to locking around

Author: ktoso

stdlib/public/Concurrency/TaskGroup.cpp

@@ -273,7 +273,6 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
     }
   };
 
-protected:
 #if SWIFT_STDLIB_SINGLE_THREADED_CONCURRENCY || SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
   // Synchronization is simple here. In a single threaded mode, all swift tasks
   // run on a single thread so no coordination is needed. In a task-to-thread
@@ -294,6 +293,7 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
   void unlock() const { mutex_.unlock(); }
 #endif
 
+protected:
   /// Used for queue management, counting number of waiting and ready tasks
   std::atomic<uint64_t> status;
 
@@ -351,6 +351,8 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
 
   /// Attempt to park the `waitingTask` in the waiting queue.
   ///
+  /// Must be called while holding `group.lock`.
+  ///
   /// If unable to complete the waiting task immediately (with an readily
   /// available completed task), either returns an `PollStatus::Empty`
   /// result if it is known that there are no pending tasks in the group,
@@ -526,7 +528,7 @@ struct TaskGroupStatus {
   ///     TaskGroupStatus{ C:{cancelled} W:{waiting task} R:{ready tasks} P:{pending tasks} {binary repr} }
   /// If discarding results:
   ///     TaskGroupStatus{ C:{cancelled} W:{waiting task} P:{pending tasks} {binary repr} }
-  std::string to_string(const TaskGroupBase* _Nonnull group) {
+  std::string to_string(const TaskGroupBase* group) {
     std::string str;
     str.append("TaskGroupStatus{ ");
     str.append("C:"); // cancelled
@@ -553,7 +555,7 @@ struct TaskGroupStatus {
 bool TaskGroupBase::statusCompletePendingReadyWaiting(TaskGroupStatus &old) {
   return status.compare_exchange_strong(
       old.status, old.completingPendingReadyWaiting(this).status,
-      /*success*/ std::memory_order_relaxed,
+      /*success*/ std::memory_order_release,
       /*failure*/ std::memory_order_relaxed);
 }
 
@@ -1148,7 +1150,7 @@ void AccumulatingTaskGroup::offer(AsyncTask *completedTask, AsyncContext *contex
     hadErrorResult = true;
   }
 
-  SWIFT_TASK_GROUP_DEBUG_LOG(this, "ready: %d, pending: %u",
+  SWIFT_TASK_GROUP_DEBUG_LOG(this, "ready: %d, pending: %llu",
                        assumed.readyTasks(this), assumed.pendingTasks(this));
 
   // ==== a) has waiting task, so let us complete it right away
@@ -1201,7 +1203,8 @@ void DiscardingTaskGroup::offer(AsyncTask *completedTask, AsyncContext *context)
 
   /// If we're the last task we've been waiting for, and there is a waiting task on the group
   bool lastPendingTaskAndWaitingTask =
-      assumed.pendingTasks(this) == 1 && assumed.hasWaitingTask();
+      assumed.pendingTasks(this) == 1 &&
+      assumed.hasWaitingTask();
 
   // Immediately decrement the pending count.
   // We can do this, since in this mode there is no ready count to keep track of,
@@ -1297,6 +1300,8 @@ void TaskGroupBase::resumeWaitingTask(
       if (statusCompletePendingReadyWaiting(assumed)) {
         // Run the task.
         auto result = PollResult::get(completedTask, hadErrorResult);
+        SWIFT_TASK_GROUP_DEBUG_LOG(this, "resume waiting DONE, task = %p, complete with = %p, status  = %s",
+                                   waitingTask, completedTask, statusString().c_str());
 
         // Remove the child from the task group's running tasks list.
         // The parent task isn't currently running (we're about to wake
@@ -1652,6 +1657,7 @@ static void swift_taskGroup_waitAllImpl(
   waitingTask->ResumeContext = rawContext;
 
   auto group = asBaseImpl(_group);
+  group->lock();
   PollResult polled = group->waitAll(bodyError, waitingTask);
 
   auto context = static_cast<TaskFutureWaitAsyncContext *>(rawContext);
@@ -1666,8 +1672,8 @@ static void swift_taskGroup_waitAllImpl(
 
   switch (polled.status) {
     case PollStatus::MustWait: {
-      SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl MustWait, pending tasks exist, waiting task = %p",
-                                 waitingTask);
+      group->unlock();
+
       // The waiting task has been queued on the channel,
       // there were pending tasks so it will be woken up eventually.
 #ifdef __ARM_ARCH_7K__
@@ -1679,6 +1685,8 @@ static void swift_taskGroup_waitAllImpl(
     }
 
     case PollStatus::Error: {
+      group->unlock();
+
       SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl Error, waiting task = %p, body error = %p, status:%s",
                                  waitingTask, bodyError, group->statusString().c_str());
 #if SWIFT_TASK_GROUP_BODY_THROWN_ERROR_WINS
@@ -1703,6 +1711,8 @@ static void swift_taskGroup_waitAllImpl(
 
     case PollStatus::Empty:
     case PollStatus::Success: {
+      group->unlock();
+
       /// Anything else than a "MustWait" can be treated as a successful poll.
       /// Only if there are in flight pending tasks do we need to wait after all.
       SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl %s, waiting task = %p, status:%s",
@@ -1721,9 +1731,8 @@ static void swift_taskGroup_waitAllImpl(
   }
 }
 
+/// Must be called while holding `group.lock`.
 PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask) {
-  lock(); // TODO: remove group lock, and use status for synchronization
-
   SWIFT_TASK_GROUP_DEBUG_LOG(this, "waitAll, bodyError = %p, status = %s", bodyError, statusString().c_str());
   PollResult result = PollResult::getEmpty(this->successType);
   result.status = PollStatus::Empty;
@@ -1739,8 +1748,7 @@ PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask)
   // otherwise we don't modify the status
   auto assumed = statusLoadAcquire();
 
-  SWIFT_TASK_GROUP_DEBUG_LOG(this, "waitAll, LOAD STATUS, status = %s",
-                             assumed.to_string(this).c_str());
+  SWIFT_TASK_GROUP_DEBUG_LOG(this, "waitAll, status = %s", assumed.to_string(this).c_str());
 
   // ==== 1) may be able to bail out early if no tasks are pending -------------
   if (assumed.isEmpty(this)) {
@@ -1758,16 +1766,13 @@ PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask)
           result.status = PollStatus::Error;
         }
       } // else, we're definitely Empty
-
-      unlock();
       return result;
     }
 
     SWIFT_TASK_GROUP_DEBUG_LOG(this, "group is empty, no pending tasks, status = %s", assumed.to_string(this).c_str());
     // No tasks in flight, we know no tasks were submitted before this poll
     // was issued, and if we parked here we'd potentially never be woken up.
     // Bail out and return `nil` from `group.next()`.
-    unlock();
     return result;
   }
 
@@ -1795,7 +1800,6 @@ PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask)
         /*success*/ std::memory_order_release,
         /*failure*/ std::memory_order_acquire)) {
       statusMarkWaitingAssumeRelease();
-      unlock(); // TODO: remove fragment lock, and use status for synchronization
 #if SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL
       // The logic here is paired with the logic in TaskGroupBase::offer. Once
       // we run the