[TaskGroup] fix missing retain in scheduling next() immediately on offer

Author: ktoso

include/swift/ABI/TaskGroup.h

@@ -37,9 +37,9 @@ namespace swift {
 
   class TaskGroup {
   public:
-    /// Describes the status of the channel.
+    /// Describes the status of the group.
     enum class ReadyStatus : uintptr_t {
-        /// The channel is empty, no tasks are pending.
+        /// The task group is empty, no tasks are pending.
         /// Return immediately, there is no point in suspending.
         ///
         /// The storage is not accessible.
@@ -54,10 +54,10 @@ namespace swift {
     };
 
     enum class PollStatus : uintptr_t {
-        /// The channel is known to be empty and we can immediately return nil.
+        /// The group is known to be empty and we can immediately return nil.
         Empty = 0,
 
-        /// The task has been enqueued to the channels wait queue.
+        /// The task has been enqueued to the groups wait queue.
         MustWait = 1,
 
         /// The task has completed with result (of type \c resultType).
@@ -68,7 +68,7 @@ namespace swift {
         Error = 3,
     };
 
-    /// The result of waiting on a Channel (TaskGroup).
+    /// The result of waiting on the TaskGroup.
     struct PollResult {
         PollStatus status; // TODO: pack it into storage pointer or not worth it?
 
@@ -81,10 +81,10 @@ namespace swift {
         /// object itself.
         OpaqueValue *storage;
 
-        /// Optional, the completed task that was polled out of the ready queue.
+        /// The completed task, if necessary to keep alive until consumed by next().
         ///
         /// # Important: swift_release
-        /// If if a task is returned here, the task MUST be swift_release'd
+        /// If if a task is returned here, the task MUST be swift_released
         /// once we are done with it, to balance out the retain made before
         /// when the task was enqueued into the ready queue to keep it alive
         /// until a next() call eventually picks it up.
@@ -96,8 +96,7 @@ namespace swift {
               status == PollStatus::Empty;
         }
 
-        static PollResult get(AsyncTask *asyncTask, bool hadErrorResult,
-                                   bool needsSwiftRelease) {
+        static PollResult get(AsyncTask *asyncTask, bool hadErrorResult) {
           auto fragment = asyncTask->futureFragment();
           return PollResult{
               /*status*/ hadErrorResult ?
@@ -106,14 +105,12 @@ namespace swift {
               /*storage*/ hadErrorResult ?
                   reinterpret_cast<OpaqueValue *>(fragment->getError()) :
                   fragment->getStoragePtr(),
-              /*task*/ needsSwiftRelease ?
-                  asyncTask :
-                  nullptr
+              /*task*/ asyncTask
           };
         }
     };
 
-    /// An item within the message queue of a channel.
+    /// An item within the message queue of a group.
     struct ReadyQueueItem {
         /// Mask used for the low status bits in a message queue item.
         static const uintptr_t statusMask = 0x03;
@@ -263,7 +260,7 @@ namespace swift {
     /// we have to keep this pointer here so we know which record to remove then.
     TaskGroupTaskStatusRecord* Record;
 
-    /// Queue containing completed tasks offered into this channel.
+    /// Queue containing completed tasks offered into this group.
     ///
     /// The low bits contain the status, the rest of the pointer is the
     /// AsyncTask.
@@ -284,7 +281,7 @@ namespace swift {
 //          readyQueue(ReadyQueueItem::get(ReadyStatus::Empty, nullptr)),
           waitQueue(nullptr) {}
 
-    /// Destroy the storage associated with the channel.
+    /// Destroy the storage associated with the group.
     void destroy(AsyncTask *task);
 
     bool isEmpty() {
@@ -373,8 +370,11 @@ namespace swift {
     }
 
 
-    /// Offer result of a task into this channel.
-    /// The value is enqueued at the end of the channel.
+    /// Offer result of a task into this task group.
+    ///
+    /// If possible, and an existing task is already waiting on next(), this will
+    /// schedule it immediately. If not, the result is enqueued and will be picked
+    /// up whenever a task calls next() the next time.
     void offer(AsyncTask *completed, AsyncContext *context, ExecutorRef executor);
 
     /// Attempt to dequeue ready tasks and complete the waitingTask.

include/swift/Runtime/Concurrency.h

@@ -161,11 +161,8 @@ SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 TaskFutureWaitThrowingSignature::FunctionType
 swift_task_future_wait_throwing;
 
-//using TaskFutureWaitSignature =
-//  AsyncSignature<void(AsyncTask *, OpaqueValue *), /*throws*/ false>;
-
-using TaskGroupFutureWaitSignature =
-  AsyncSignature<void(AsyncTask *, TaskGroup *, OpaqueValue *), /*throws*/ true>; // TODO: is this correct?
+using TaskGroupFutureWaitThrowingSignature =
+  AsyncSignature<void(AsyncTask *, TaskGroup *, Metadata *), /*throws*/ true>;
 
 /// Wait for a readyQueue of a Channel to become non empty.
 ///
@@ -178,7 +175,7 @@ using TaskGroupFutureWaitSignature =
 /// ) async -> T
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
-TaskGroupFutureWaitSignature::FunctionType
+TaskGroupFutureWaitThrowingSignature::FunctionType
 swift_task_group_wait_next_throwing;
 
 /// Create a new `TaskGroup` using the task's allocator.
@@ -194,23 +191,6 @@ swift_task_group_wait_next_throwing;
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 swift::TaskGroup* swift_task_group_create(AsyncTask *task);
 
-///// Attach the task group to the passed in (current) task that runs it.
-/////
-///// Returns the allocated and attached record that must be removed from the task,
-///// when the group exits.
-/////
-///// Its Swift signature is
-/////
-///// \code
-///// func swift_task_group_attach(
-/////     group: Builtin.NativeObject
-/////     to task: Builtin.NativeObject
-///// ) -> UnsafeRawPointer
-///// \endcode
-//SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-//TaskGroupTaskStatusRecord*
-//swift_task_group_attach(TaskGroup *group, AsyncTask *task);
-
 /// Attach a child task to the parent task's task group record.
 ///
 /// Its Swift signature is
@@ -226,20 +206,6 @@ SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_group_attachChild(TaskGroup *group,
                                   AsyncTask *parent, AsyncTask *child);
 
-///// Detach the group from its parent and remove the passed in record.
-/////
-///// Its Swift signature is
-/////
-///// \code
-///// func swift_task_group_detach(
-/////     group: UnsafeRawPointer,
-/////     from task: Builtin.NativeObject
-///// ) -> UnsafeRawPointer
-///// \endcode
-//SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-//void
-//swift_task_group_detach(TaskGroupTaskStatusRecord *record, AsyncTask *task);
-
 /// Its Swift signature is
 ///
 /// \code

stdlib/public/Concurrency/Task.cpp

@@ -115,7 +115,6 @@ void AsyncTask::completeFuture(AsyncContext *context, ExecutorRef executor) {
     // then we must offer into the parent group that we completed,
     // so it may `next()` poll completed child tasks in completion order.
     auto group = groupChildFragment()->getGroup();
-    assert(group);
     group->offer(this, context, executor);
   }
 
@@ -147,7 +146,6 @@ void AsyncTask::completeFuture(AsyncContext *context, ExecutorRef executor) {
 SWIFT_CC(swift)
 static void destroyTask(SWIFT_CONTEXT HeapObject *obj) {
   auto task = static_cast<AsyncTask*>(obj);
-
   // For a future, destroy the result.
   if (task->isFuture()) {
     task->futureFragment()->destroy();

stdlib/public/Concurrency/TaskGroup.cpp

@@ -146,6 +146,16 @@ void TaskGroup::offer(AsyncTask *completedTask, AsyncContext *context,
   assert(completedTask->hasGroupChildFragment());
   assert(completedTask->groupChildFragment()->getGroup() == this);
 
+  // We retain the completed task, because we will either:
+  // - (a) schedule the waiter to resume on the next() that it is waiting on, or
+  // - (b) will need to store this task until the group task enters next() and
+  //       picks up this task.
+  // either way, there is some time between us returning here, and the `completeTask`
+  // issuing a swift_release on this very task. We need to keep it alive until
+  // we have the chance to poll it from the queue (via the waiter task entering
+  // calling next()).
+  swift_retain(completedTask);
+
   mutex.lock(); // TODO: remove fragment lock, and use status for synchronization
 
   // Immediately increment ready count and acquire the status
@@ -179,11 +189,10 @@ void TaskGroup::offer(AsyncTask *completedTask, AsyncContext *context,
           /*failure*/ std::memory_order_acquire) &&
           statusCompletePendingReadyWaiting(assumed)) {
         // Run the task.
-        auto result = PollResult::get(
-            completedTask, hadErrorResult, /*needsRelease*/ false);
+        auto result = PollResult::get(completedTask, hadErrorResult);
 
         mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
-        // swift::runTaskWithPollResult(waitingTask, completingExecutor, result);
+
         auto waitingContext =
           static_cast<TaskFutureWaitAsyncContext *>(
             waitingTask->ResumeContext);
@@ -193,8 +202,6 @@ void TaskGroup::offer(AsyncTask *completedTask, AsyncContext *context,
         swift_task_enqueueGlobal(waitingTask);
         return;
       } // else, try again
-
-      assert(false && "why should this have to try again ever?"); // FIXME
     }
   }
 
@@ -234,9 +241,8 @@ void swift::swift_task_group_wait_next_throwing(
   auto context = static_cast<TaskFutureWaitAsyncContext *>(rawContext);
   auto task = context->task;
   auto group = context->group;
-fprintf(stderr, "[%s:%d](%s) group: %d\n", __FILE_NAME__, __LINE__, __FUNCTION__, group);
-
-assert(group && "swift_task_group_wait_next_throwing was passed context without group!");
+  assert(waitingTask == task && "attempted to wait on group.next() from other task, which is illegal!");
+  assert(group && "swift_task_group_wait_next_throwing was passed context without group!");
 
   TaskGroup::PollResult polled = group->poll(waitingTask);
   switch (polled.status) {

stdlib/public/Concurrency/TaskGroup.swift

@@ -220,29 +220,6 @@ extension Task {
       #endif
 
       return try await _taskGroupWaitNext(waitingTask: _task, group: _group)
-//      let rawResult = await _taskGroupWaitNext(waitingTask: _task, group: _group)
-//
-//      // FIXME: we are now in the parent task again, and can safely modify the
-//      //        records again; we should remove this specific child from the
-//      //        group->record
-//
-//      if rawResult.hadErrorResult {
-//        // Throw the result on error.
-//        let error = unsafeBitCast(rawResult.storage, to: Error.self)
-//        throw error
-//      }
-//
-//      guard let storage = rawResult.storage else {
-//        // The group was empty, return nil
-//        return nil
-//      }
-//
-//      // Take the value on success.
-//      let storagePtr =
-//        storage.bindMemory(to: TaskResult.self, capacity: 1)
-//      let value = UnsafeMutablePointer<TaskResult>(mutating: storagePtr).pointee
-//
-//      return value
     }
 
     /// Query whether the group has any remaining tasks.

stdlib/public/Concurrency/TaskPrivate.h

@@ -81,8 +81,9 @@ class TaskFutureWaitAsyncContext : public AsyncContext {
   // Arguments.
   AsyncTask *task;
 
-  // Only in swift_task_group_next
+  // Only in swift_task_group_wait_next_throwing.
   TaskGroup *group;
+  // Only in swift_task_group_wait_next_throwing.
   const Metadata *successType;
 
   using AsyncContext::AsyncContext;

test/Concurrency/Runtime/async_task_equals_hashCode.swift

@@ -3,7 +3,6 @@
 // REQUIRES: executable_test
 // REQUIRES: concurrency
 
-import Dispatch
 #if canImport(Darwin)
 import Darwin
 #elseif canImport(Glibc)

test/Concurrency/Runtime/async_taskgroup_is_asyncsequence.swift

@@ -1,13 +1,7 @@
-// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency %import-libdispatch -parse-as-library) | %FileCheck %s --dump-input=always
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input=always
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: libdispatch
-// FIXME: unlock on other OSes
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
-
-import Dispatch
 
 #if canImport(Darwin)
 import Darwin

test/Concurrency/Runtime/async_taskgroup_throw_recover.swift

@@ -1,11 +1,7 @@
-// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency %import-libdispatch -parse-as-library) | %FileCheck %s --dump-input=always
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input=always
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: libdispatch
-// FIXME: unlock on other OSes
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
 
 struct Boom: Error {}
 struct IgnoredBoom: Error {}

test/Concurrency/Runtime/async_taskgroup_throw_rethrow.swift

@@ -1,10 +1,7 @@
-// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency %import-libdispatch -parse-as-library) | %FileCheck %s
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: libdispatch
-
-import Dispatch
 
 struct Boom: Error {}
 struct IgnoredBoom: Error {}
@@ -29,11 +26,11 @@ func test_taskGroup_throws_rethrows() async {
         throw error
       }
 
-      fatalError("should have thrown")
+      print("should have thrown")
+      return 0
     }
 
-      print("got: \(got)")
-      fatalError("Expected error to be thrown, but got: \(got)")
+    print("Expected error to be thrown, but got: \(got)")
   } catch {
     // CHECK: rethrown: Boom()
     print("rethrown: \(error)")