Offering body error must be done while holding lock

Author: ktoso

stdlib/public/Concurrency/TaskGroup.cpp

@@ -51,15 +51,24 @@
 using namespace swift;
 
 #if 0
-#define SWIFT_TASK_GROUP_DEBUG_LOG(group, fmt, ...) \
+#define SWIFT_TASK_GROUP_DEBUG_LOG(group, fmt, ...)                     \
 fprintf(stderr, "[%#lx] [%s:%d][group(%p%s)] (%s) " fmt "\n",           \
       (unsigned long)Thread::current().platformThreadId(),              \
       __FILE__, __LINE__,                                               \
       group, group->isDiscardingResults() ? ",discardResults" : "",     \
       __FUNCTION__,                                                     \
       __VA_ARGS__)
+
+#define SWIFT_TASK_GROUP_DEBUG_LOG_0(group, fmt, ...)                   \
+fprintf(stderr, "[%#lx] [%s:%d][group(%p)] (%s) " fmt "\n",             \
+      (unsigned long)Thread::current().platformThreadId(),              \
+      __FILE__, __LINE__,                                               \
+      group,                                                            \
+      __FUNCTION__,                                                     \
+      __VA_ARGS__)
 #else
 #define SWIFT_TASK_GROUP_DEBUG_LOG(group, fmt, ...) (void)0
+#define SWIFT_TASK_GROUP_DEBUG_LOG_0(group, fmt, ...) (void)0
 #endif
 
 using FutureFragment = AsyncTask::FutureFragment;
@@ -354,7 +363,11 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
   /// There can be only at-most-one waiting task on a group at any given time,
   /// and the waiting task is expected to be the parent task in which the group
   /// body is running.
-  PollResult waitAll(AsyncTask *waitingTask);
+  ///
+  /// \param bodyError error thrown by the body of a with...TaskGroup method
+  /// \param waitingTask the task waiting on the group
+  /// \return how the waiting task should be handled, e.g. must wait or can be completed immediately
+  PollResult waitAll(SwiftError* bodyError, AsyncTask *waitingTask);
 
   // Enqueue the completed task onto ready queue if there are no waiting tasks yet
   virtual void enqueueCompletedTask(AsyncTask *completedTask, bool hadErrorResult) = 0;
@@ -411,6 +424,15 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
   virtual TaskGroupStatus statusAddPendingTaskRelaxed(bool unconditionally) = 0;
 };
 
+static std::string to_string(TaskGroupBase::PollStatus status) {
+  switch (status) {
+    case TaskGroupBase::PollStatus::Empty: return "Empty";
+    case TaskGroupBase::PollStatus::MustWait: return "MustWait";
+    case TaskGroupBase::PollStatus::Success: return "Success";
+    case TaskGroupBase::PollStatus::Error: return "Error";
+  }
+}
+
 /// The status of a task group.
 ///
 /// Its exact structure depends on the type of group, and therefore a group must be passed to operations
@@ -855,8 +877,8 @@ static void swift_taskGroup_initializeWithFlagsImpl(size_t rawGroupFlags,
                                                     TaskGroup *group, const Metadata *T) {
 
   TaskGroupFlags groupFlags(rawGroupFlags);
-  SWIFT_TASK_DEBUG_LOG("group(%p) create; flags: isDiscardingResults=%d",
-                       group, groupFlags.isDiscardResults());
+  SWIFT_TASK_GROUP_DEBUG_LOG_0(group, "create group; flags: isDiscardingResults=%d",
+                               groupFlags.isDiscardResults());
 
   TaskGroupBase *impl;
   if (groupFlags.isDiscardResults()) {
@@ -1618,7 +1640,7 @@ static void swift_taskGroup_waitAllImpl(
   waitingTask->ResumeContext = rawContext;
 
   auto group = asBaseImpl(_group);
-  PollResult polled = group->waitAll(waitingTask);
+  PollResult polled = group->waitAll(bodyError, waitingTask);
 
   auto context = static_cast<TaskFutureWaitAsyncContext *>(rawContext);
   context->ResumeParent =
@@ -1627,23 +1649,13 @@ static void swift_taskGroup_waitAllImpl(
   context->errorResult = nullptr;
   context->successResultPointer = resultPointer;
 
-  SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl, waiting task = %p, bodyError = %p, status:%s",
-                       waitingTask, bodyError, group->statusString().c_str());
+  SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl, waiting task = %p, bodyError = %p, status:%s, polled.status = %s",
+                       waitingTask, bodyError, group->statusString().c_str(), to_string(polled.status).c_str());
 
   switch (polled.status) {
     case PollStatus::MustWait:
-    SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAll MustWait, pending tasks exist, waiting task = %p",
+    SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl MustWait, pending tasks exist, waiting task = %p",
                                waitingTask);
-      if (bodyError && group->isDiscardingResults()) {
-        auto discardingGroup = asDiscardingImpl(_group);
-        bool storedBodyError = discardingGroup->offerBodyError(bodyError);
-        if (storedBodyError) {
-          SWIFT_TASK_GROUP_DEBUG_LOG(
-              group, "waitAll, stored error thrown by with...Group body, error = %p",
-              bodyError);
-        }
-      }
-
       // The waiting task has been queued on the channel,
       // there were pending tasks so it will be woken up eventually.
 #ifdef __ARM_ARCH_7K__
@@ -1654,7 +1666,7 @@ static void swift_taskGroup_waitAllImpl(
 #endif /* __ARM_ARCH_7K__ */
 
     case PollStatus::Error:
-      SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAll found error, waiting task = %p, body error = %p, status:%s",
+      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
       if (bodyError) {
@@ -1676,17 +1688,10 @@ static void swift_taskGroup_waitAllImpl(
       return waitingTask->runInFullyEstablishedContext();
 
     case PollStatus::Empty:
-      /// 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, "waitAll %s, waiting task = %p, status:%s",
-                                 polled.status == TaskGroupBase::PollStatus::Empty ? "empty" : "success",
-                                 waitingTask, group->statusString().c_str());
-
-
     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, "waitAll %s, waiting task = %p, status:%s",
+      SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAllImpl %s, waiting task = %p, status:%s",
                                  polled.status == TaskGroupBase::PollStatus::Empty ? "empty" : "success",
                                  waitingTask, group->statusString().c_str());
 
@@ -1717,10 +1722,13 @@ bool DiscardingTaskGroup::offerBodyError(SwiftError* _Nonnull bodyError) {
   return true;
 }
 
-PollResult TaskGroupBase::waitAll(AsyncTask *waitingTask) {
+/// 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.
+PollResult TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask) {
   lock(); // TODO: remove group lock, and use status for synchronization
-
-  SWIFT_TASK_GROUP_DEBUG_LOG(this, "waitAll, status = %s", statusString().c_str());
+  
+  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;
   result.storage = nullptr;
@@ -1764,6 +1772,18 @@ PollResult TaskGroupBase::waitAll(AsyncTask *waitingTask) {
   }
 
   // ==== 2) Add to wait queue -------------------------------------------------
+
+  // ---- 2.1) Discarding task group may need to story the bodyError before we park
+  if (bodyError && isDiscardingResults()) {
+    auto discardingGroup = asDiscardingImpl(this);
+    bool storedBodyError = discardingGroup->offerBodyError(bodyError);
+    if (storedBodyError) {
+      SWIFT_TASK_GROUP_DEBUG_LOG(
+          this, "waitAll, stored error thrown by with...Group body, error = %p",
+          bodyError);
+    }
+  }
+
   auto waitHead = waitQueue.load(std::memory_order_acquire);
   _swift_tsan_release(static_cast<Job *>(waitingTask));
   while (true) {