prepare for throwing

ktoso · ktoso · commit ebd2e994ddcb · 2022-12-20T08:46:50.000+09:00
diff --git a/stdlib/public/BackDeployConcurrency/TaskGroup.cpp b/stdlib/public/BackDeployConcurrency/TaskGroup.cpp
@@ -566,7 +566,7 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
   }
 
   case PollStatus::Empty: {
-    // Initialize the result as a nil Optional<Success>.
+    // Initialize the result as a .none Optional<Success>.
     const Metadata *successType = result.successType;
     OpaqueValue *destPtr = context->successResultPointer;
     successType->vw_storeEnumTagSinglePayload(destPtr, 1, 1);
diff --git a/stdlib/public/Concurrency/TaskGroup.cpp b/stdlib/public/Concurrency/TaskGroup.cpp
@@ -54,11 +54,13 @@ using namespace swift;
 /*************************** TASK GROUP ***************************************/
 /******************************************************************************/
 
-#if SWIFT_TASK_DEBUG_LOG_ENABLED
-#define SWIFT_TASK_GROUP_DEBUG_LOG(group, fmt, ...)                                   \
-  SWIFT_TASK_DEBUG_LOG("group(%p%s) " fmt,                                            \
-          group, group->isDiscardingResults() ? ",discardResults" : "",               \
-          __VA_ARGS__)
+#if 1
+#define SWIFT_TASK_GROUP_DEBUG_LOG(group, fmt, ...) \
+fprintf(stderr, "[%#lx] [%s:%d](%s) group(%p%s) " fmt "\n",             \
+      (unsigned long)Thread::current().platformThreadId(),              \
+      __FILE__, __LINE__, __FUNCTION__,                                 \
+      group, group->isDiscardingResults() ? ",discardResults" : "",     \
+      __VA_ARGS__)
 #else
 #define SWIFT_TASK_GROUP_DEBUG_LOG(group, fmt, ...) (void)0
 #endif
@@ -212,13 +214,13 @@ class TaskGroupImpl: public TaskGroupTaskStatusRecord {
     }
 
     unsigned int readyTasks(const TaskGroupImpl* _Nonnull group) {
-      assert(group && group->isAccumulatingResults()
+      assert(group->isAccumulatingResults()
              && "attempted to check ready tasks on group that does not accumulate results!");
       return (status & maskReady) >> 31;
     }
 
-    uint64_t pendingTasks(const TaskGroupImpl* _Nullable group) {
-      if (group && group->isAccumulatingResults()) {
+    uint64_t pendingTasks(const TaskGroupImpl* _Nonnull group) {
+      if (group->isAccumulatingResults()) {
         return (status & maskAccumulatingPending);
       } else {
         return (status & maskDiscardingPending);
@@ -260,17 +262,16 @@ class TaskGroupImpl: public TaskGroupTaskStatusRecord {
     ///     GroupStatus{ C:{cancelled} W:{waiting task} R:{ready tasks} P:{pending tasks} {binary repr} }
     /// If discarding results:
     ///     GroupStatus{ C:{cancelled} W:{waiting task} P:{pending tasks} {binary repr} }
-    std::string to_string(const TaskGroupImpl* _Nullable group) {
+    std::string to_string(const TaskGroupImpl* _Nonnull group) {
       std::string str;
       str.append("GroupStatus{ ");
       str.append("C:"); // cancelled
-      str.append(isCancelled() ? "y " : "n ");
-      str.append("W:"); // has waiting task
+      str.append(isCancelled() ? "y" : "n");
+      str.append(" W:"); // has waiting task
       str.append(hasWaitingTask() ? "y" : "n");
       if (group && group->isAccumulatingResults()) {
         str.append(" R:"); // ready
         str.append(std::to_string(readyTasks(group)));
-        str.append(" ");
       }
       str.append(" P:"); // pending
       str.append(std::to_string(pendingTasks(group)));
@@ -411,7 +412,7 @@ class TaskGroupImpl: public TaskGroupTaskStatusRecord {
   GroupStatus statusMarkWaitingAssumeAcquire() {
     fprintf(stderr, "[%s:%d](%s) statusMarkWaitingAssumeAcquire, load....\n", __FILE_NAME__, __LINE__, __FUNCTION__);
     fprintf(stderr, "[%s:%d](%s) statusMarkWaitingAssumeAcquire = %s\n", __FILE_NAME__, __LINE__, __FUNCTION__,
-            statusLoadRelaxed().to_string(nullptr).c_str());
+            statusLoadRelaxed().to_string(this).c_str());
     auto old = status.fetch_or(GroupStatus::waiting, std::memory_order_acquire);
     return GroupStatus{old | GroupStatus::waiting};
   }
@@ -693,12 +694,16 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
     // Initialize the result as an Optional<Success>.
     const Metadata *successType = result.successType;
     OpaqueValue *destPtr = context->successResultPointer;
-    successType->vw_storeEnumTagSinglePayload(destPtr, 1, 1);
+    // TODO: figure out a way to try to optimistically take the
+    // value out of the finished task's future, if there are no
+    // remaining references to it.
+    successType->vw_initializeWithCopy(destPtr, result.storage);
+    successType->vw_storeEnumTagSinglePayload(destPtr, 0, 1);
     return;
   }
 
   case PollStatus::Empty: {
-    // Initialize the result as an Optional<Success>.
+    // Initialize the result as a .none Optional<Success>.
     const Metadata *successType = result.successType;
     OpaqueValue *destPtr = context->successResultPointer;
     successType->vw_storeEnumTagSinglePayload(destPtr, 1, 1);
@@ -709,37 +714,10 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
 
 static void fillGroupNextNilResult(TaskFutureWaitAsyncContext *context,
                                    PollResult result) {
-//  /// Fill in the result value
-//  switch (result.status) {
-//  case PollStatus::MustWait:
-//    assert(false && "filling a waiting status?");
-//    return;
-//
-//  case PollStatus::Error: {
-//    assert(false && "cannot have errors");
-//    return;
-//  }
-//
-//  case PollStatus::Success: {
-//    // Initialize the result as an Optional<Void>.
-//    const Metadata *successType = result.successType;
-//    OpaqueValue *destPtr = context->successResultPointer;
-//    // TODO: figure out a way to try to optimistically take the
-//    // value out of the finished task's future, if there are no
-//    // remaining references to it.
-//    successType->vw_initializeWithCopy(destPtr, result.storage);
-//    successType->vw_storeEnumTagSinglePayload(destPtr, 0, 1);
-//    return;
-//  }
-//
-//  case PollStatus::Empty: {
-    // Initialize the result as a nil Optional<Success>.
-    const Metadata *successType = result.successType;
-    OpaqueValue *destPtr = context->successResultPointer;
-    successType->vw_storeEnumTagSinglePayload(destPtr, 1, 1);
-    return;
-//  }
-//  }
+  // Initialize the result as a .none Optional<Success>.
+  const Metadata *successType = result.successType;
+  OpaqueValue *destPtr = context->successResultPointer;
+  successType->vw_storeEnumTagSinglePayload(destPtr, 1, 1);
 }
 
 // TaskGroup is locked upon entry and exit
@@ -774,7 +752,7 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
   assert(completedTask->hasChildFragment());
   assert(completedTask->hasGroupChildFragment());
   assert(completedTask->groupChildFragment()->getGroup() == asAbstract(this));
-  SWIFT_TASK_GROUP_DEBUG_LOG(this, "offer, completedTask:%p , status:%s", completedTask, statusLoadRelaxed().to_string(this).c_str());
+  SWIFT_TASK_GROUP_DEBUG_LOG(this, "ENTER offer, completedTask:%p , status:%s", completedTask, statusLoadRelaxed().to_string(this).c_str());
 
   // The current ownership convention is that we are *not* given ownership
   // of a retain on completedTask; we're called from the task completion
@@ -802,19 +780,19 @@ void TaskGroupImpl::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_DEBUG_LOG("group(%p) discard result, was pending:%llu",
-                         this, assumed.pendingTasks(this));
+    SWIFT_TASK_GROUP_DEBUG_LOG(this, "discard result, was pending:%llu",
+                               assumed.pendingTasks(this));
 
     // 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 (assumed.pendingTasks(this) == 1 && assumed.hasWaitingTask()) {
-      SWIFT_TASK_DEBUG_LOG("group(%p) offer, discardResults, offered last pending task, resume waiting task:%p",
-                           this, waitQueue.load(std::memory_order_relaxed));
+      SWIFT_TASK_GROUP_DEBUG_LOG(this, "offer, discardResults, offered last pending task, resume waiting task:%p",
+                           waitQueue.load(std::memory_order_relaxed));
       resumeWaitingTask(completedTask, assumed, /*hadErrorResult=*/false);
     } else {
       auto afterComplete = statusCompletePendingAssumeRelease(this);
-      SWIFT_TASK_DEBUG_LOG("group(%p) offer, discardResults, either more pending tasks, or no waiting task, status:%s", this,
+      SWIFT_TASK_GROUP_DEBUG_LOG(this, "offer, either more pending tasks, or no waiting task, status:%s",
                            afterComplete.to_string(this).c_str());
       _swift_taskGroup_detachChild(asAbstract(this), completedTask);
     }
@@ -823,8 +801,8 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
     return;
   }
 
-  SWIFT_TASK_DEBUG_LOG("group(%p), ready: %d, pending: %llu",
-                       this, assumed.readyTasks(this), assumed.pendingTasks(this));
+  SWIFT_TASK_GROUP_DEBUG_LOG(this, "ready: %d, pending: %llu",
+                       assumed.readyTasks(this), assumed.pendingTasks(this));
 
   auto asyncContextPrefix = reinterpret_cast<FutureAsyncContextPrefix *>(
       reinterpret_cast<char *>(context) - sizeof(FutureAsyncContextPrefix));
@@ -859,8 +837,8 @@ void TaskGroupImpl::resumeWaitingTask(
     TaskGroupImpl::GroupStatus &assumed,
     bool hadErrorResult) {
   auto waitingTask = waitQueue.load(std::memory_order_acquire);
-  SWIFT_TASK_DEBUG_LOG("group(%p) resume waiting task = %p, complete with = %p",
-                       this, waitingTask, completedTask);
+  SWIFT_TASK_GROUP_DEBUG_LOG(this, "resume waiting task = %p, complete with = %p",
+                       waitingTask, completedTask);
   while (true) {
     // ==== a) run waiting task directly -------------------------------------
     assert(assumed.hasWaitingTask());
@@ -1186,7 +1164,6 @@ static void swift_taskGroup_waitAllImpl(
     TaskGroup *_group,
     ThrowingTaskFutureWaitContinuationFunction *resumeFunction,
     AsyncContext *rawContext) {
-  fprintf(stderr, "[%s:%d](%s) group(%p) WAIT ALL IMPL\n", __FILE_NAME__, __LINE__, __FUNCTION__, _group);
   auto waitingTask = swift_task_getCurrent();
   waitingTask->ResumeTask = task_group_wait_resume_adapter;
   waitingTask->ResumeContext = rawContext;
@@ -1199,8 +1176,8 @@ static void swift_taskGroup_waitAllImpl(
   context->successResultPointer = resultPointer;
 
   auto group = asImpl(_group);
-  SWIFT_TASK_DEBUG_LOG("group(%p) waitAll, waiting task = %p, status:%s",
-                       group, waitingTask, group->statusLoadRelaxed().to_string(group).c_str());
+  SWIFT_TASK_GROUP_DEBUG_LOG(group, "waitAll, waiting task = %p, status:%s",
+                       waitingTask, group->statusLoadRelaxed().to_string(group).c_str());
 
   PollResult polled = group->waitAll(waitingTask);
   switch (polled.status) {
diff --git a/stdlib/public/Concurrency/TaskGroup.swift b/stdlib/public/Concurrency/TaskGroup.swift
@@ -105,15 +105,12 @@ public func withTaskGroup<ChildTaskResult, GroupResult>(
 
   let _group = Builtin.createTaskGroupWithFlags(flags, ChildTaskResult.self)
   var group = TaskGroup<ChildTaskResult>(group: _group)
+  defer { Builtin.destroyTaskGroup(_group) }
 
   // Run the withTaskGroup body.
   let result = await body(&group)
 
-  let x: ChildTaskResult? = try? await _taskGroupWaitAll(group: _group) // FIXME: should not need to throw
-  assert(x == nil)
-  // await group.awaitAllRemainingTasks() // FIXME: use this instead!
-
-  Builtin.destroyTaskGroup(_group)
+  let _: ChildTaskResult? = try? await _taskGroupWaitAll(group: _group) // try!-safe, cannot throw since this is a non throwing group
   return result
 }
 
@@ -201,14 +198,14 @@ public func withThrowingTaskGroup<ChildTaskResult, GroupResult>(
     // Run the withTaskGroup body.
     let result = try await body(&group)
 
-    await group.awaitAllRemainingTasks()
+    _ = try? await group.awaitAllRemainingTasks()
     Builtin.destroyTaskGroup(_group)
 
     return result
   } catch {
     group.cancelAll()
 
-    await group.awaitAllRemainingTasks()
+    _ = try? await group.awaitAllRemainingTasks() // discard errors
     Builtin.destroyTaskGroup(_group)
 
     throw error
@@ -219,6 +216,45 @@ public func withThrowingTaskGroup<ChildTaskResult, GroupResult>(
   #endif
 }
 
+@available(SwiftStdlib 5.8, *)
+@_unsafeInheritExecutor
+@inlinable
+public func withThrowingTaskGroup<ChildTaskResult, GroupResult>(
+  of childTaskResultType: ChildTaskResult.Type,
+  returning returnType: GroupResult.Type = GroupResult.self,
+  discardResults: Bool,
+  body: (inout ThrowingTaskGroup<ChildTaskResult, Error>) async throws -> GroupResult
+) async throws -> GroupResult {
+  let flags = taskGroupCreateFlags(
+          discardResults: discardResults
+  )
+
+  let _group = Builtin.createTaskGroupWithFlags(flags, ChildTaskResult.self)
+  var group = ThrowingTaskGroup<ChildTaskResult, Error>(group: _group)
+
+  do {
+    // Run the withTaskGroup body.
+    let result = try await body(&group)
+
+    try await group.awaitAllRemainingTasks()
+    Builtin.destroyTaskGroup(_group)
+
+    return result
+  } catch {
+    group.cancelAll()
+
+    do {
+      try await group.awaitAllRemainingTasks()
+      Builtin.destroyTaskGroup(_group)
+    } catch {
+      Builtin.destroyTaskGroup(_group)
+      throw error
+    }
+
+    throw error
+  }
+}
+
 /// A group that contains dynamically created child tasks.
 ///
 /// To create a task group,
@@ -478,8 +514,7 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
     /// implementation.
     if #available(SwiftStdlib 5.8, *) {
       if isDiscardingResults {
-        let x: ChildTaskResult? = try? await _taskGroupWaitAll(group: _group) // FIXME: should not need to throw
-        assert(x == nil)
+        let _: ChildTaskResult? = try! await _taskGroupWaitAll(group: _group) // try!-safe, cannot throw, not throwing group
         return
       }
     }
@@ -514,7 +549,7 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
   ///
   /// - SeeAlso: ``waitForAll`` which waits for all pending tasks to complete.
   @available(SwiftStdlib 5.8, *)
-  internal var isDiscardingResults: Bool {
+  public var isDiscardingResults: Bool {
     _taskGroupIsDiscardingResults(group: _group) // TODO: test this
   }
 
@@ -606,7 +641,20 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
 
   /// Await all the remaining tasks on this group.
   @usableFromInline
-  internal mutating func awaitAllRemainingTasks() async {
+  internal mutating func awaitAllRemainingTasks() async throws {
+    /// Since 5.8, we implement "wait for all pending tasks to complete"
+    /// in the runtime, in order to be able to handle the discard-results
+    /// implementation.
+    if #available(SwiftStdlib 5.8, *) {
+      if isDiscardingResults {
+        let _: ChildTaskResult? = try await _taskGroupWaitAll(group: _group) // if any of the tasks throws, this will "rethrow" here
+        return
+      }
+    }
+
+    // Old implementation just consumes the group as normal end-user code would.
+    // As it does not have to account for the `discardResults` mode.
+    // The first error encountered while draining tasks is rethrown.
     while true {
       do {
         guard let _ = try await next() else {
@@ -618,13 +666,13 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
 
   @usableFromInline
   internal mutating func _waitForAll() async throws {
-    while let _ = try? await next() { }
+    try await self.awaitAllRemainingTasks()
   }
 
   /// Wait for all of the group's remaining tasks to complete.
   @_alwaysEmitIntoClient
   public mutating func waitForAll() async throws {
-    while let _ = try await next() { }
+    try await self.awaitAllRemainingTasks()
   }
 
 #if !SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY
@@ -904,6 +952,18 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
     _taskGroupIsEmpty(_group)
   }
 
+  /// A Boolean value that indicates whether the group has been created in `discardResults` mode.
+  ///
+  /// If true, the group will not be accumulating results and the `next()` method will always
+  /// return `nil` immediately.
+  ///
+  /// - SeeAlso: ``waitForAll`` which waits for all pending tasks to complete.
+  @available(SwiftStdlib 5.8, *)
+  public var isDiscardingResults: Bool {
+    _taskGroupIsDiscardingResults(group: _group) // TODO: test this
+  }
+
+
   /// Cancel all of the remaining tasks in the group.
   ///
   /// After cancellation,
@@ -1143,6 +1203,7 @@ func _taskHasTaskGroupStatusRecord() -> Bool
 /// Always returns `nil`.
 @available(SwiftStdlib 5.8, *)
 @usableFromInline
+@discardableResult
 @_silgen_name("swift_taskGroup_waitAll")
 func _taskGroupWaitAll<T>(group: Builtin.RawPointer) async throws -> T?
 
diff --git a/stdlib/public/Concurrency/TaskPrivate.h b/stdlib/public/Concurrency/TaskPrivate.h
@@ -38,7 +38,7 @@ namespace swift {
 
 // Set to 1 to enable helpful debug spew to stderr
 // If this is enabled, tests with `swift_task_debug_log` requirement can run.
-#if 1
+#if 0
 #define SWIFT_TASK_DEBUG_LOG_ENABLED 1
 #define SWIFT_TASK_DEBUG_LOG(fmt, ...)                                         \
   fprintf(stderr, "[%#lx] [%s:%d](%s) " fmt "\n",                              \
diff --git a/test/Concurrency/Runtime/async_taskgroup_throw_rethrow.swift b/test/Concurrency/Runtime/async_taskgroup_throw_rethrow.swift

Original file line number	Diff line number	Diff line change
`@@ -566,7 +566,7 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,`
`566`	`566`	`}`
`567`	`567`
`568`	`568`	`case PollStatus::Empty: {`
`569`		`- // Initialize the result as a nil Optional<Success>.`
	`569`	`+ // Initialize the result as a .none Optional<Success>.`
`570`	`570`	`const Metadata *successType = result.successType;`
`571`	`571`	`OpaqueValue *destPtr = context->successResultPointer;`
`572`	`572`	`successType->vw_storeEnumTagSinglePayload(destPtr, 1, 1);`