mark waiting status only when we actually MustWait, while holding group lock

Author: ktoso

stdlib/public/Concurrency/TaskGroup.cpp

@@ -378,6 +378,7 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
   // ==== Status manipulation -------------------------------------------------
 
   TaskGroupStatus statusLoadRelaxed() const;
+  TaskGroupStatus statusLoadAcquire() const;
 
   std::string statusString() const;
 
@@ -409,6 +410,10 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
   /// Remove waiting status bit.
   TaskGroupStatus statusRemoveWaitingRelease();
 
+  /// Mark the waiting status bit.
+  /// A waiting task MUST have been already enqueued in the `waitQueue`.
+  TaskGroupStatus statusMarkWaitingAssumeRelease();
+
   /// Cancels the group and returns true if was already cancelled before.
   /// After this function returns, the group is guaranteed to be cancelled.
   ///
@@ -561,6 +566,10 @@ TaskGroupStatus TaskGroupBase::statusLoadRelaxed() const {
   return TaskGroupStatus{status.load(std::memory_order_relaxed)};
 }
 
+TaskGroupStatus TaskGroupBase::statusLoadAcquire() const {
+  return TaskGroupStatus{status.load(std::memory_order_acquire)};
+}
+
 std::string TaskGroupBase::statusString() const {
   return statusLoadRelaxed().to_string(this);
 }
@@ -580,6 +589,12 @@ TaskGroupStatus TaskGroupBase::statusMarkWaitingAssumeAcquire() {
   return TaskGroupStatus{old | TaskGroupStatus::waiting};
 }
 
+TaskGroupStatus TaskGroupBase::statusMarkWaitingAssumeRelease() {
+  auto old = status.fetch_or(TaskGroupStatus::waiting,
+                             std::memory_order_release);
+  return TaskGroupStatus{old | TaskGroupStatus::waiting};
+}
+
 TaskGroupStatus TaskGroupBase::statusRemoveWaitingRelease() {
   auto old = status.fetch_and(~TaskGroupStatus::waiting,
                               std::memory_order_release);
@@ -702,18 +717,6 @@ class DiscardingTaskGroup: public TaskGroupBase {
     return true;
   }
 
-  /// Returns *assumed* new status, including the just performed +1.
-  TaskGroupStatus statusMarkWaitingAssumeAcquire() {
-    auto old = status.fetch_or(TaskGroupStatus::waiting, std::memory_order_acquire);
-    return TaskGroupStatus{old | TaskGroupStatus::waiting};
-  }
-
-  TaskGroupStatus statusRemoveWaitingRelease() {
-    auto old = status.fetch_and(~TaskGroupStatus::waiting,
-                                std::memory_order_release);
-    return TaskGroupStatus{old};
-  }
-
   /// Returns *assumed* new status.
   TaskGroupStatus statusAddReadyAssumeAcquire(const DiscardingTaskGroup *group) {
     assert(group->isDiscardingResults());
@@ -1203,8 +1206,8 @@ void DiscardingTaskGroup::offer(AsyncTask *completedTask, AsyncContext *context)
   // Immediately decrement the pending count.
   // We can do this, since in this mode there is no ready count to keep track of,
   // and we immediately discard the result.
-  SWIFT_TASK_GROUP_DEBUG_LOG(this, "discard result, hadError:%d, was pending:%llu",
-                             hadErrorResult, assumed.pendingTasks(this));
+  SWIFT_TASK_GROUP_DEBUG_LOG(this, "discard result, hadError:%d, was pending:%llu, status = %s",
+                             hadErrorResult, assumed.pendingTasks(this), assumed.to_string(this).c_str());
   // If this was the last pending task, and there is a waiting task (from waitAll),
   // we must resume the task; but not otherwise. There cannot be any waiters on next()
   // while we're discarding results.
@@ -1662,9 +1665,9 @@ static void swift_taskGroup_waitAllImpl(
                        waitingTask, bodyError, group->statusString().c_str(), to_string(polled.status).c_str());
 
   switch (polled.status) {
-    case PollStatus::MustWait:
-    SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl MustWait, pending tasks exist, waiting task = %p",
-                               waitingTask);
+    case PollStatus::MustWait: {
+      SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl MustWait, pending tasks exist, waiting task = %p",
+                                 waitingTask);
       // The waiting task has been queued on the channel,
       // there were pending tasks so it will be woken up eventually.
 #ifdef __ARM_ARCH_7K__
@@ -1673,8 +1676,9 @@ static void swift_taskGroup_waitAllImpl(
 #else /* __ARM_ARCH_7K__ */
       return;
 #endif /* __ARM_ARCH_7K__ */
+    }
 
-    case PollStatus::Error:
+    case PollStatus::Error: {
       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
@@ -1695,9 +1699,10 @@ static void swift_taskGroup_waitAllImpl(
       }
 
       return waitingTask->runInFullyEstablishedContext();
+    }
 
     case PollStatus::Empty:
-    case PollStatus::Success:
+    case PollStatus::Success: {
       /// 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",
@@ -1712,12 +1717,11 @@ static void swift_taskGroup_waitAllImpl(
       }
 
       return waitingTask->runInFullyEstablishedContext();
+    }
   }
 }
 
-/// Must be called while holding the `taskGroup.lock`!
-/// This is because the discarding task group still has some follow-up operations that must
-/// be performed atomically after this operation sometimes, so we cannot unlock inside `waitAll` itself.
+/// Caller must mark the `waiting` status bit when MustWait is returned.
 PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask) {
   lock(); // TODO: remove group lock, and use status for synchronization
 
@@ -1732,7 +1736,12 @@ PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask)
   bool haveRunOneChildTaskInline = false;
 
   reevaluate_if_TaskGroup_has_results:;
-  auto assumed = statusMarkWaitingAssumeAcquire();
+  // Paired with a release when marking Waiting,
+  // 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());
 
   // ==== 1) may be able to bail out early if no tasks are pending -------------
   if (assumed.isEmpty(this)) {
@@ -1759,7 +1768,6 @@ PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask)
     // 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()`.
-    statusRemoveWaitingRelease();
     unlock();
     return result;
   }
@@ -1787,6 +1795,7 @@ PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask)
         waitHead, 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