complete impl, except body throws

Author: ktoso

stdlib/public/Concurrency/TaskGroup.cpp

@@ -577,7 +577,6 @@ TaskGroupTaskStatusRecord * TaskGroup::getTaskRecord() {
 // Initializes into the preallocated _group an actual TaskGroupImpl.
 SWIFT_CC(swift)
 static void swift_taskGroup_initializeImpl(TaskGroup *group, const Metadata *T) {
-  fprintf(stderr, "[%s:%d](%s) INITIALIZE...\n", __FILE_NAME__, __LINE__, __FUNCTION__);
   swift_taskGroup_initializeWithFlags(0, group, T);
 }
 
@@ -588,9 +587,6 @@ static void swift_taskGroup_initializeWithFlagsImpl(size_t rawGroupFlags, TaskGr
 
   TaskGroupFlags groupFlags(rawGroupFlags);
 
-  fprintf(stderr, "[%s:%d](%s) INITIALIZE FLAGS: flags.discardResults:%d\n", __FILE_NAME__, __LINE__, __FUNCTION__,
-          groupFlags.isDiscardResults());
-
   TaskGroupImpl *impl = ::new(group)
     TaskGroupImpl(T, groupFlags.isDiscardResults());
   auto record = impl->getTaskRecord();
@@ -809,23 +805,28 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
     // and we immediately discard the result.
     SWIFT_TASK_GROUP_DEBUG_LOG(this, "discard result, hadError:%d, was pending:%llu",
                                hadErrorResult, assumed.pendingTasks(this));
-    if (!lastPendingTaskAndWaitingTask) {
-      // we're not able to immediately complete a waitingTask with this task, so we may have to store it...
-      if (hadErrorResult && readyQueue.isEmpty()) {
-        // a discardResults throwing task group must retain the FIRST error it encounters.
-        SWIFT_TASK_GROUP_DEBUG_LOG(this, "offer error, completedTask:%p", completedTask);
-        enqueueCompletedTask(completedTask, /*hadErrorResult=*/hadErrorResult);
-      }
-    } // else, no need to store the task, as we'll immediately complete the waitingTask using it.
-
     // 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.
     if (lastPendingTaskAndWaitingTask) {
+      ReadyQueueItem item;
+      bool dequeuedErrorItem = readyQueue.dequeue(item);
       SWIFT_TASK_GROUP_DEBUG_LOG(this, "offer, offered last pending task, resume waiting task:%p",
                                  waitQueue.load(std::memory_order_relaxed));
-      resumeWaitingTask(completedTask, assumed, /*hadErrorResult=*/hadErrorResult);
+      if (dequeuedErrorItem) {
+        assert(item.getStatus() == ReadyStatus::Error && "only errors can be stored by a discarding task group, yet it wasn't an error!");
+        resumeWaitingTask(item.getTask(), assumed, /*hadErrorResult=*/true);
+      } else {
+        resumeWaitingTask(completedTask, assumed, /*hadErrorResult=*/hadErrorResult);
+      }
     } else {
+      assert(!lastPendingTaskAndWaitingTask);
+      if (hadErrorResult && readyQueue.isEmpty()) {
+        // a discardResults throwing task group must retain the FIRST error it encounters.
+        SWIFT_TASK_GROUP_DEBUG_LOG(this, "offer error, completedTask:%p", completedTask);
+        enqueueCompletedTask(completedTask, /*hadErrorResult=*/hadErrorResult);
+      } // else, we just are going to discard it.
+
       auto afterComplete = statusCompletePendingAssumeRelease();
       SWIFT_TASK_GROUP_DEBUG_LOG(this, "offer, either more pending tasks, or no waiting task, status:%s",
                                  afterComplete.to_string(this).c_str());

test/Concurrency/Runtime/async_taskgroup_throw_rethrow.swift

@@ -87,10 +87,12 @@ func test_discardingTaskGroup_automaticallyRethrows() async {
 }
 
 func test_discardingTaskGroup_automaticallyRethrowsOnlyFirst() async {
-  print("==== \(#function) ------") // CHECK_LABEL: test_discardingTaskGroup_automaticallyRethrowsOnlyFirst
+  print("==== \(#function) ------") // CHECK-LABEL: test_discardingTaskGroup_automaticallyRethrowsOnlyFirst
   do {
     let got = try await withThrowingDiscardingTaskGroup(returning: Int.self) { group in
-      group.addTask { await echo(1) }
+      group.addTask {
+        await echo(1)
+      }
       group.addTask {
         let error = Boom(id: "first, isCancelled:\(Task.isCancelled)")
         print("Throwing: \(error)")
@@ -102,7 +104,7 @@ func test_discardingTaskGroup_automaticallyRethrowsOnlyFirst() async {
         do {
           try await Task.sleep(until: .now + .seconds(120), clock: .continuous)
         } catch {
-          print("Throwing: \(error)")
+          print("Awoken, throwing: \(error)")
           throw error
         }
       }
@@ -112,46 +114,13 @@ func test_discardingTaskGroup_automaticallyRethrowsOnlyFirst() async {
     print("Expected error to be thrown, but got: \(got)")
   } catch {
     // CHECK: Throwing: Boom(id: "first, isCancelled:false
-    // CHECK: Throwing: CancellationError()
+    // CHECK: Awoken, throwing: CancellationError()
     // and only then the re-throw happens:
     // CHECK: rethrown: Boom(id: "first
     print("rethrown: \(error)")
   }
 }
 
-func test_discardingTaskGroup_automaticallyRethrowsOnlyFirstIncludingGroupBody() async {
-  print("==== \(#function) ------") // CHECK_LABEL: test_discardingTaskGroup_automaticallyRethrowsOnlyFirstIncludingGroupBody
-  do {
-    try await withThrowingDiscardingTaskGroup(returning: Int.self) { group in
-      group.addTask { await echo(1) }
-      group.addTask {
-        try await Task.sleep(until: .now + .seconds(1), clock: .continuous)
-        let error = Boom(id: "first, isCancelled:\(Task.isCancelled)")
-        print("Throwing: \(error)")
-        throw error
-      }
-      group.addTask {
-        // we wait "forever" but since the group will get cancelled after
-        // the first error, this will be woken up and throw a cancellation
-        try await Task.sleep(until: .now + .seconds(20), clock: .continuous)
-      }
-
-      let bodyError = Boom(id: "body, isCancelled:\(group.isCancelled)")
-      print("Throwing: \(bodyError)")
-      throw bodyError
-    }
-
-    print("Expected error to be thrown")
-  } catch {
-    // CHECK: Throwing: Boom(id: "body, isCancelled:false
-    // CHECK: Throwing: Boom(id: "first, isCancelled:true
-    // CHECK: Throwing: Boom(id: "second, isCancelled:true
-    // and only then the re-throw happens:
-    // CHECK: rethrown: Boom(id: "body
-    print("rethrown: \(error)")
-  }
-}
-
 @available(SwiftStdlib 5.1, *)
 @main struct Main {
   static func main() async {