[Concurrency] TaskGroup children remove their records as they complete

If we didn't do this (and we didn't), the tasks get released as we
perform the next() impl, and move the value from the ready task to the
waiting task. Then, the ready task gets destroyed.

But as the task group exists, it performs a cancelAll() and that
iterates over all records. Those records were not removed previously
(!!!) which meant we were pointing at now deallocated tasks.

Previously this worked because we didn't deallocate the tasks, so they
leaked, but we didn't crash. With the memory leak fixed, this began to
crash since we'd attempt to cancel already destroyed tasks.

Solution:
- Remove task records whenever they complete a waiting task.
- This can ONLY be done by the "group owning task" itself, becuause
  the contract of ONLY this task being allowed to modify records. o
  It MUST NOT be done by the completing tasks as they complete, as it
  would race with the owning task modifying this linked list of child
  tasks in the group record.

(cherry picked from commit f336404)

Author: ktoso

include/swift/ABI/TaskStatus.h

@@ -182,7 +182,6 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
 
   /// Attach the passed in `child` task to this group.
   void attachChild(AsyncTask *child) {
-    assert(child->groupChildFragment());
     assert(child->hasGroupChildFragment());
     assert(child->groupChildFragment()->getGroup() == getGroup());
 
@@ -211,6 +210,33 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
     cur->childFragment()->setNextChild(child);
   }
 
+  void detachChild(AsyncTask *child) {
+    assert(child && "cannot remove a null child from group");
+    if (FirstChild == child) {
+      FirstChild = getNextChildTask(child);
+      return;
+    }
+
+    AsyncTask *prev = FirstChild;
+    // Remove the child from the linked list, i.e.:
+    //     prev -> afterPrev -> afterChild
+    //                 ==
+    //               child   -> afterChild
+    // Becomes:
+    //     prev --------------> afterChild
+    while (prev) {
+      auto afterPrev = getNextChildTask(prev);
+
+      if (afterPrev == child) {
+        auto afterChild = getNextChildTask(child);
+        prev->childFragment()->setNextChild(afterChild);
+        return;
+      }
+
+      prev = afterPrev;
+    }
+  }
+
   static AsyncTask *getNextChildTask(AsyncTask *task) {
     return task->childFragment()->getNextChild();
   }

stdlib/public/Concurrency/TaskGroup.cpp

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -278,7 +278,7 @@ class TaskGroupImpl: public TaskGroupTaskStatusRecord {
 
 private:
 
-//    // TODO: move to lockless via the status atomic
+  // TODO: move to lockless via the status atomic (make readyQueue an mpsc_queue_t<ReadyQueueItem>)
   mutable std::mutex mutex;
 
   /// Used for queue management, counting number of waiting and ready tasks
@@ -289,7 +289,6 @@ class TaskGroupImpl: public TaskGroupTaskStatusRecord {
   /// The low bits contain the status, the rest of the pointer is the
   /// AsyncTask.
   NaiveQueue<ReadyQueueItem> readyQueue;
-//     mpsc_queue_t<ReadyQueueItem> readyQueue; // TODO: can we get away with an MPSC queue here once actor executors land?
 
   /// Single waiting `AsyncTask` currently waiting on `group.next()`,
   /// or `nullptr` if no task is currently waiting.
@@ -304,10 +303,8 @@ class TaskGroupImpl: public TaskGroupTaskStatusRecord {
     : TaskGroupTaskStatusRecord(),
       status(GroupStatus::initial().status),
       readyQueue(),
-//          readyQueue(ReadyQueueItem::get(ReadyStatus::Empty, nullptr)),
       waitQueue(nullptr), successType(T) {}
 
-
   TaskGroupTaskStatusRecord *getTaskRecord() {
     return reinterpret_cast<TaskGroupTaskStatusRecord *>(this);
   }
@@ -471,11 +468,18 @@ static void swift_taskGroup_initializeImpl(TaskGroup *group, const Metadata *T)
 
 // =============================================================================
 // ==== add / attachChild ------------------------------------------------------
+
 SWIFT_CC(swift)
 static void swift_taskGroup_attachChildImpl(TaskGroup *group,
                                             AsyncTask *child) {
+  SWIFT_TASK_DEBUG_LOG("attach child task = %p to group = %p\n",
+          child, group);
+
+  // The counterpart of this (detachChild) is performed by the group itself,
+  // when it offers the completed (child) task's value to a waiting task -
+  // during the implementation of `await group.next()`.
   auto groupRecord = asImpl(group)->getTaskRecord();
-  return groupRecord->attachChild(child);
+  groupRecord->attachChild(child);
 }
 
 // =============================================================================
@@ -505,8 +509,7 @@ bool TaskGroup::isCancelled() {
 }
 
 static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
-                                PollResult result,
-                                bool releaseResultRetainedTask) {
+                                PollResult result) {
   /// Fill in the result value
   switch (result.status) {
   case PollStatus::MustWait:
@@ -515,7 +518,7 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
 
   case PollStatus::Error: {
     context->fillWithError(reinterpret_cast<SwiftError *>(result.storage));
-    break;
+    return;
   }
 
   case PollStatus::Success: {
@@ -527,28 +530,17 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
     // remaining references to it.
     successType->vw_initializeWithCopy(destPtr, result.storage);
     successType->vw_storeEnumTagSinglePayload(destPtr, 0, 1);
-    break;
+    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);
-    break;
+    return;
   }
   }
-
-  // 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) {
@@ -557,7 +549,7 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
   assert(completedTask->hasChildFragment());
   assert(completedTask->hasGroupChildFragment());
   assert(completedTask->groupChildFragment()->getGroup() == asAbstract(this));
-  SWIFT_TASK_DEBUG_LOG("offer task %p to group %p\n", completedTask, group);
+  SWIFT_TASK_DEBUG_LOG("offer task %p to group %p", completedTask, this);
 
   mutex.lock(); // TODO: remove fragment lock, and use status for synchronization
 
@@ -603,7 +595,9 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
             static_cast<TaskFutureWaitAsyncContext *>(
                 waitingTask->ResumeContext);
 
-        fillGroupNextResult(waitingContext, result, /*release*/false);
+        fillGroupNextResult(waitingContext, result);
+        detachChild(result.retainedTask);
+
         _swift_tsan_acquire(static_cast<Job *>(waitingTask));
 
         // TODO: allow the caller to suggest an executor
@@ -621,11 +615,12 @@ 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",
+  SWIFT_TASK_DEBUG_LOG("group has no waiting tasks, RETAIN and 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);
+
   auto readyItem = ReadyQueueItem::get(
       hadErrorResult ? ReadyStatus::Error : ReadyStatus::Success,
       completedTask
@@ -691,7 +686,7 @@ 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",
+    SWIFT_TASK_DEBUG_LOG("poll group = %p, no ready tasks, waiting task = %p",
                          group, waitingTask);
     // The waiting task has been queued on the channel,
     // there were pending tasks so it will be woken up eventually.
@@ -705,15 +700,22 @@ static void swift_taskGroup_wait_next_throwingImpl(
   case PollStatus::Empty:
   case PollStatus::Error:
   case PollStatus::Success:
-    SWIFT_TASK_DEBUG_LOG("poll group = %p, task = %p, ready task available = %p\n",
+    SWIFT_TASK_DEBUG_LOG("poll group = %p, task = %p, ready task available = %p",
                          group, waitingTask, polled.retainedTask);
-    fillGroupNextResult(context, polled, /*release*/true);
+    fillGroupNextResult(context, polled);
+    if (auto completedTask = polled.retainedTask) {
+      // it would be null for PollStatus::Empty, then we don't need to release
+      group->detachChild(polled.retainedTask);
+      swift_release(polled.retainedTask);
+    }
+
     return waitingTask->runInFullyEstablishedContext();
   }
 }
 
 PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
   mutex.lock(); // TODO: remove group lock, and use status for synchronization
+  SWIFT_TASK_DEBUG_LOG("poll group = %p", this);
   auto assumed = statusMarkWaitingAssumeAcquire();
 
   PollResult result;
@@ -723,6 +725,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
 
   // ==== 1) bail out early if no tasks are pending ----------------------------
   if (assumed.isEmpty()) {
+    SWIFT_TASK_DEBUG_LOG("poll group = %p, group is empty, no pending tasks", this);
     // No tasks in flight, we know no tasks were submitted before this poll
     // was issued, and if we parked here we'd potentially never be woken up.
     // Bail out and return `nil` from `group.next()`.
@@ -740,6 +743,9 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
 
   // ==== 2) Ready task was polled, return with it immediately -----------------
   if (assumed.readyTasks()) {
+    SWIFT_TASK_DEBUG_LOG("poll group = %p, group has ready tasks = %d",
+                         this, assumed.readyTasks());
+
     auto assumedStatus = assumed.status;
     auto newStatus = TaskGroupImpl::GroupStatus{assumedStatus};
     if (status.compare_exchange_weak(
@@ -843,13 +849,16 @@ static void swift_taskGroup_cancelAllImpl(TaskGroup *group) {
 }
 
 bool TaskGroupImpl::cancelAll() {
+  SWIFT_TASK_DEBUG_LOG("cancel all tasks in group = %p", this);
+
   // store the cancelled bit
   auto old = statusCancel();
   if (old.isCancelled()) {
     // already was cancelled previously, nothing to do?
     return false;
   }
 
+  // FIXME: must also remove the records!!!!
   // cancel all existing tasks within the group
   swift_task_cancel_group_child_tasks(asAbstract(this));
   return true;

stdlib/public/Concurrency/TaskPrivate.h

@@ -33,6 +33,7 @@ namespace swift {
 //#define SWIFT_TASK_PRINTF_DEBUG 1
 
 // Set to 1 to enable helpful debug spew to stderr
+// If this is enabled, tests with `swift_task_debug_log` requirement can run.
 #if 0
 #define SWIFT_TASK_DEBUG_LOG(fmt, ...)                                         \
   fprintf(stderr, "[%lu] [%s:%d](%s) " fmt "\n",                               \

stdlib/public/Concurrency/TaskStatus.cpp

@@ -354,6 +354,8 @@ static bool swift_task_tryAddStatusRecordImpl(TaskStatusRecord *newRecord) {
 SWIFT_CC(swift)
 static bool swift_task_removeStatusRecordImpl(TaskStatusRecord *record) {
   auto task = swift_task_getCurrent();
+  SWIFT_TASK_DEBUG_LOG("remove status record = %p, from current task = %p",
+                       record, task);
 
   // Load the current state.
   auto &status = task->_private().Status;
@@ -454,6 +456,8 @@ static ChildTaskStatusRecord*
 swift_task_attachChildImpl(AsyncTask *child) {
   void *allocation = malloc(sizeof(swift::ChildTaskStatusRecord));
   auto record = new (allocation) swift::ChildTaskStatusRecord(child);
+  SWIFT_TASK_DEBUG_LOG("attach child task = %p, record = %p, to current task = %p",
+                       child, record, swift_task_getCurrent());
   swift_task_addStatusRecord(record);
   return record;
 }
@@ -548,6 +552,7 @@ static void performGroupCancellationAction(TaskStatusRecord *record) {
 
 SWIFT_CC(swift)
 static void swift_task_cancelImpl(AsyncTask *task) {
+  SWIFT_TASK_DEBUG_LOG("cancel task = %p", task);
   Optional<StatusRecordLockRecord> recordLockRecord;
 
   // Acquire the status record lock.

test/Concurrency/Runtime/async_task_withUnsafeCurrentTask.swift

@@ -0,0 +1,39 @@
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking -parse-as-library) 2>&1 | %FileCheck %s --dump-input=always
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: swift_task_debug_log
+
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+#if os(Linux)
+import Glibc
+#elseif os(Windows)
+import MSVCRT
+#else
+import Darwin
+#endif
+
+func test_withUnsafeCurrentTask() async {
+  // The task we're running in ("main")
+  // CHECK: creating task [[MAIN_TASK:0x.*]] with parent 0x0
+
+  // CHECK: creating task [[TASK:0x.*]] with parent 0x0
+  let t = Task.detached {
+    withUnsafeCurrentTask { task in
+      fputs("OK: \(task!)", stderr)
+    }
+    fputs("DONE", stderr)
+  }
+
+  // CHECK: OK: UnsafeCurrentTask(_task: (Opaque Value))
+  // CHECK: DONE
+  // CHECK: destroy task [[TASK]]
+  await t.value
+}
+
+@main struct Main {
+  static func main() async {
+    await test_withUnsafeCurrentTask()
+  }
+}

test/Concurrency/Runtime/async_taskgroup_dontLeakTasks.swift

@@ -1,9 +1,10 @@
-// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking -parse-as-library) | %FileCheck %s
+// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking -parse-as-library) 2>&1 | %FileCheck %s --dump-input=always
 // REQUIRES: executable_test
 // REQUIRES: concurrency
+// REQUIRES: swift_task_debug_log
+
 // UNSUPPORTED: use_os_stdlib
 // UNSUPPORTED: back_deployment_runtime
-// UNSUPPORTED: linux
 
 #if os(Linux)
 import Glibc
@@ -13,10 +14,16 @@ import MSVCRT
 import Darwin
 #endif
 
-@available(SwiftStdlib 5.5, *)
 func test_taskGroup_next() async {
+  // CHECK: creating task [[MAIN_TASK:0x.*]] with parent 0x0
+  // CHECK: creating task [[GROUP_TASK_1:0x.*]] with parent [[MAIN_TASK]]
+  // CHECK: creating task [[GROUP_TASK_2:0x.*]] with parent [[MAIN_TASK]]
+  // CHECK: creating task [[GROUP_TASK_3:0x.*]] with parent [[MAIN_TASK]]
+  // CHECK: creating task [[GROUP_TASK_4:0x.*]] with parent [[MAIN_TASK]]
+  // CHECK: creating task [[GROUP_TASK_5:0x.*]] with parent [[MAIN_TASK]]
+
   _ = await withTaskGroup(of: Int.self, returning: Int.self) { group in
-    for n in 0..<100 {
+    for n in 0..<5 {
       group.spawn {
         return n
       }
@@ -30,12 +37,18 @@ func test_taskGroup_next() async {
 
     return sum
   }
-
-  // CHECK: result with group.next(): 100
-  print("result with group.next(): \(100)")
+  // as we exit the group, it must be guaranteed that its child tasks were destroyed
+  //
+  // NOTE: there is no great way to express "any of GROUP_TASK_n",
+  //       so we just check that 5 tasks were destroyed
+  //
+  // CHECK: destroy task [[DESTROY_GROUP_TASK_1:0x.*]]
+  // CHECK: destroy task [[DESTROY_GROUP_TASK_2:0x.*]]
+  // CHECK: destroy task [[DESTROY_GROUP_TASK_3:0x.*]]
+  // CHECK: destroy task [[DESTROY_GROUP_TASK_4:0x.*]]
+  // CHECK: destroy task [[DESTROY_GROUP_TASK_5:0x.*]]
 }
 
-@available(SwiftStdlib 5.5, *)
 @main struct Main {
   static func main() async {
     await test_taskGroup_next()

test/Concurrency/Runtime/async_taskgroup_throw_recover.swift

@@ -41,7 +41,15 @@ func test_taskGroup_throws() async {
         return 3
       }
 
-      guard let third = try! await group.next() else {
+      switch await group.nextResult() {
+      case .success(let third):
+        print("task group returning normally: \(third)")
+        return third
+
+      case .failure(let error):
+        fatalError("got an erroneous third result: \(error)")
+
+      case .none:
         print("task group failed to get 3")
         return 0
       }