[Concurrency] Fix group child tasks not being released

The proper handling of task group child tasks is that:
- if it completes a waiting task immediately, we don't need to retain it
  - we just move the value to the waiting task and can destroy the task
- if we need to store the ready task and wait for a waiting task (for a
  task that hits `await group.next()`) then we need to retain the ready
  task.
  - as the waiting task arrives, we move the value from the ready task
    to the waiting task, and swift_release the ready task -- it will now
    be destroyed safely.

Author: ktoso

stdlib/public/Concurrency/Task.cpp

@@ -144,7 +144,7 @@ void NullaryContinuationJob::process(Job *_job) {
 void AsyncTask::completeFuture(AsyncContext *context) {
   using Status = FutureFragment::Status;
   using WaitQueueItem = FutureFragment::WaitQueueItem;
-
+  SWIFT_TASK_DEBUG_LOG("complete future = %p", this);
   assert(isFuture());
   auto fragment = futureFragment();
 
@@ -230,7 +230,6 @@ AsyncTask::~AsyncTask() {
 SWIFT_CC(swift)
 static void destroyTask(SWIFT_CONTEXT HeapObject *obj) {
   auto task = static_cast<AsyncTask*>(obj);
-
   task->~AsyncTask();
 
   // The task execution itself should always hold a reference to it, so

stdlib/public/Concurrency/TaskGroup.cpp

@@ -490,16 +490,8 @@ void TaskGroupImpl::destroy() {
   // First, remove the group from the task and deallocate the record
   swift_task_removeStatusRecord(getTaskRecord());
 
-  mutex.lock(); // TODO: remove lock, and use status for synchronization
-  // Release all ready tasks which are kept retained, the group destroyed,
-  // so no other task will ever await on them anymore;
-  ReadyQueueItem item;
-  bool taskDequeued = readyQueue.dequeue(item);
-  while (taskDequeued) {
-    swift_release(item.getTask());
-    taskDequeued = readyQueue.dequeue(item);
-  }
-  mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
+  // By the time we call destroy, all tasks inside the group must have been
+  // awaited on already; We handle this on the swift side.
 }
 
 // =============================================================================
@@ -514,16 +506,18 @@ bool TaskGroup::isCancelled() {
 }
 
 static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
-                                PollResult result) {
+                                PollResult result,
+                                bool releaseResultRetainedTask) {
   /// Fill in the result value
   switch (result.status) {
   case PollStatus::MustWait:
     assert(false && "filling a waiting status?");
     return;
 
-  case PollStatus::Error:
-    context->fillWithError(reinterpret_cast<SwiftError*>(result.storage));
-    return;
+  case PollStatus::Error: {
+    context->fillWithError(reinterpret_cast<SwiftError *>(result.storage));
+    break;
+  }
 
   case PollStatus::Success: {
     // Initialize the result as an Optional<Success>.
@@ -534,17 +528,28 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
     // remaining references to it.
     successType->vw_initializeWithCopy(destPtr, result.storage);
     successType->vw_storeEnumTagSinglePayload(destPtr, 0, 1);
-    return;
+    break;
   }
 
   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;
+    break;
   }
   }
+
+  // We only release if asked to; This is because this function is called in two
+  // cases "immediately":
+  // a) when a completed task arrives and a waiting one existed then we don't
+  //    need to retain the completed task at all, thus we also don't release it.
+  // b) when the task was stored in the readyQueue it was retained. As a
+  //    waitingTask arrives we will fill-in with the value from the retained
+  //    task. In this situation we must release the ready task, to allow it to
+  //    be destroyed.
+  if (releaseResultRetainedTask)
+    swift_release(result.retainedTask);
 }
 
 void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
@@ -553,16 +558,7 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
   assert(completedTask->hasChildFragment());
   assert(completedTask->hasGroupChildFragment());
   assert(completedTask->groupChildFragment()->getGroup() == asAbstract(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);
+  SWIFT_TASK_DEBUG_LOG("offer task %p to group %p\n", completedTask, group);
 
   mutex.lock(); // TODO: remove fragment lock, and use status for synchronization
 
@@ -586,6 +582,8 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
   // ==== a) has waiting task, so let us complete it right away
   if (assumed.hasWaitingTask()) {
     auto waitingTask = waitQueue.load(std::memory_order_acquire);
+    SWIFT_TASK_DEBUG_LOG("group has waiting task = %p, complete with = %p",
+                         waitingTask, completedTask);
     while (true) {
       // ==== a) run waiting task directly -------------------------------------
       assert(assumed.hasWaitingTask());
@@ -606,15 +604,16 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
             static_cast<TaskFutureWaitAsyncContext *>(
                 waitingTask->ResumeContext);
 
-        fillGroupNextResult(waitingContext, result);
-
+        fillGroupNextResult(waitingContext, result, /*release*/false);
         _swift_tsan_acquire(static_cast<Job *>(waitingTask));
 
         // TODO: allow the caller to suggest an executor
         swift_task_enqueueGlobal(waitingTask);
         return;
       } // else, try again
     }
+
+    llvm_unreachable("should have enqueued and returned.");
   }
 
   // ==== b) enqueue completion ------------------------------------------------
@@ -623,7 +622,8 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
   // queue when a task polls during next() it will notice that we have a value
   // ready for it, and will process it immediately without suspending.
   assert(!waitQueue.load(std::memory_order_relaxed));
-
+  SWIFT_TASK_DEBUG_LOG("group has no waiting tasks, store ready task = %p",
+                       completedTask);
   // Retain the task while it is in the queue;
   // it must remain alive until the task group is alive.
   swift_retain(completedTask);
@@ -668,6 +668,7 @@ SWIFT_CC(swiftasync) static void workaround_function_swift_taskGroup_wait_next_t
 
 // =============================================================================
 // ==== group.next() implementation (wait_next and groupPoll) ------------------
+
 SWIFT_CC(swiftasync)
 static void swift_taskGroup_wait_next_throwingImpl(
     OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
@@ -691,6 +692,8 @@ static void swift_taskGroup_wait_next_throwingImpl(
   PollResult polled = group->poll(waitingTask);
   switch (polled.status) {
   case PollStatus::MustWait:
+    SWIFT_TASK_DEBUG_LOG("poll group = %p, no ready tasks, waiting task = %p\n",
+                         group, 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__
@@ -703,7 +706,9 @@ static void swift_taskGroup_wait_next_throwingImpl(
   case PollStatus::Empty:
   case PollStatus::Error:
   case PollStatus::Success:
-    fillGroupNextResult(context, polled);
+    SWIFT_TASK_DEBUG_LOG("poll group = %p, task = %p, ready task available = %p\n",
+                         group, waitingTask, polled.retainedTask);
+    fillGroupNextResult(context, polled, /*release*/true);
     return waitingTask->runInFullyEstablishedContext();
   }
 }

stdlib/public/Concurrency/TaskPrivate.h

@@ -44,7 +44,9 @@ namespace swift {
 // Set to 1 to enable helpful debug spew to stderr
 #if 0
 #define SWIFT_TASK_DEBUG_LOG(fmt, ...)                                         \
-  fprintf(stderr, "[%lu] " fmt "\n", (unsigned long)_swift_get_thread_id(),    \
+  fprintf(stderr, "[%lu] [%s:%d](%s) " fmt "\n",                               \
+          (unsigned long)_swift_get_thread_id(),                               \
+          __FILE__, __LINE__, __FUNCTION__,                                    \
           __VA_ARGS__)
 #else
 #define SWIFT_TASK_DEBUG_LOG(fmt, ...) (void)0

test/Concurrency/Runtime/async_taskgroup_asynciterator_semantics.swift

@@ -26,9 +26,11 @@ func test_taskGroup_next() async {
       do {
         while let r = try await group.next() {
           sum += r
+          print("add \(r) -> sum: \(sum)")
         }
       } catch {
         catches += 1
+        print("catch: \(catches)")
       }
     }
 

test/Concurrency/Runtime/async_taskgroup_dontLeakTasks.swift

@@ -0,0 +1,43 @@
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking -parse-as-library) | %FileCheck %s
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+// UNSUPPORTED: linux
+
+#if os(Linux)
+import Glibc
+#elseif os(Windows)
+import MSVCRT
+#else
+import Darwin
+#endif
+
+@available(SwiftStdlib 5.5, *)
+func test_taskGroup_next() async {
+  _ = await withTaskGroup(of: Int.self, returning: Int.self) { group in
+    for n in 0..<100 {
+      group.spawn {
+        return n
+      }
+    }
+    await Task.sleep(2_000_000)
+
+    var sum = 0
+    for await value in group {
+      sum += 1
+    }
+
+    return sum
+  }
+
+  // CHECK: result with group.next(): 100
+  print("result with group.next(): \(100)")
+}
+
+@available(SwiftStdlib 5.5, *)
+@main struct Main {
+  static func main() async {
+    await test_taskGroup_next()
+  }
+}

test/Concurrency/Runtime/async_taskgroup_next_on_completed.swift

@@ -15,29 +15,23 @@ func test_sum_nextOnCompleted() async {
   let numbers = [1, 2, 3, 4, 5]
   let expected = 15 // FIXME: numbers.reduce(0, +) this hangs?
 
-  let sum = try! await withTaskGroup(of: Int.self) {
-    (group) async -> Int in
+  let sum = try! await withTaskGroup(of: Int.self) { group async -> Int in
     for n in numbers {
       group.spawn {
-        () async -> Int in
         print("  complete group.spawn { \(n) }")
         return n
       }
     }
 
     // We specifically want to await on completed child tasks in this test,
     // so give them some time to complete before we hit group.next()
-    await Task.sleep(2_000_000_000)
+    try! await Task.sleep(nanoseconds: 2_000_000_000)
 
     var sum = 0
-    do {
-      while let r = try await group.next() {
-        print("next: \(r)")
-        sum += r
-        print("sum: \(sum)")
-      }
-    } catch {
-      print("ERROR: \(error)")
+    while let r = await group.next() {
+      print("next: \(r)")
+      sum += r
+      print("sum: \(sum)")
     }
 
     assert(group.isEmpty, "Group must be empty after we consumed all tasks")