[Discarding] Don't leak retained "first error" task when retaining it

Author: ktoso

stdlib/public/Concurrency/TaskGroup.cpp

@@ -228,7 +228,7 @@ class TaskGroupBase : public TaskGroupTaskStatusRecord {
                       reinterpret_cast<OpaqueValue *>(fragment->getError()) :
                       fragment->getStoragePtr(),
           /*successType=*/fragment->getResultType(),
-          /*task=*/asyncTask
+          /*retainedTask==*/asyncTask
       };
     }
 
@@ -1040,7 +1040,7 @@ static void fillGroupNextResult(TaskFutureWaitAsyncContext *context,
 
   case PollStatus::Error: {
     auto error = reinterpret_cast<SwiftError *>(result.storage);
-    fillGroupNextErrorResult(context, error);
+    fillGroupNextErrorResult(context, error); // FIXME: this specifically retains the error, but likely should not!??!!?
     return;
   }
 
@@ -1218,7 +1218,6 @@ void DiscardingTaskGroup::offer(AsyncTask *completedTask, AsyncContext *context)
   // We can do this, since in this mode there is no ready count to keep track of,
   // and we immediately discard the result.
   auto afterComplete = statusCompletePendingAssumeRelease();
-  (void) afterComplete;
   const bool alreadyDecrementedStatus = true;
   SWIFT_TASK_GROUP_DEBUG_LOG(this, "offer, complete, status afterComplete:%s", afterComplete.to_string(this).c_str());
 
@@ -1351,8 +1350,8 @@ void TaskGroupBase::resumeWaitingTask(
         // Run the task.
         auto result = PollResult::get(completedTask, hadErrorResult);
         SWIFT_TASK_GROUP_DEBUG_LOG(this,
-                                   "resume waiting DONE, task = %p, backup = %p, error:%d, complete with = %p, status = %s",
-                                   waitingTask, backup, hadErrorResult, completedTask, statusString().c_str());
+                                   "resume waiting DONE, task = %p, error:%d, complete with = %p, status = %s",
+                                   waitingTask, hadErrorResult, completedTask, statusString().c_str());
 
         auto waitingContext =
             static_cast<TaskFutureWaitAsyncContext *>(
@@ -1369,6 +1368,11 @@ void TaskGroupBase::resumeWaitingTask(
         // locks) because we know that the child task is completed and
         // we can't be holding its locks ourselves.
         _swift_taskGroup_detachChild(asAbstract(this), completedTask);
+        if (hadErrorResult) {
+          // We only used the task to keep the error in the future fragment around
+          // so now that we emitted the error and detached the task, we are free to release the task immediately.
+          swift_release(completedTask);
+        }
 
         _swift_tsan_acquire(static_cast<Job *>(waitingTask));
         // TODO: allow the caller to suggest an executor
@@ -1377,7 +1381,7 @@ void TaskGroupBase::resumeWaitingTask(
 #endif /* SWIFT_CONCURRENCY_TASK_TO_THREAD_MODEL */
     } else {
       SWIFT_TASK_GROUP_DEBUG_LOG(this, "CAS failed, task = %p, backup = %p, complete with = %p, status = %s",
-                                 waitingTask, backup, completedTask, statusString().c_str());
+                                 waitingTask, completedTask, statusString().c_str());
     }
   }
 }

test/Concurrency/Runtime/async_taskgroup_discarding_noLeaks.swift

@@ -0,0 +1,122 @@
+// RUN: %target-run-simple-leaks-swift( -Xfrontend -disable-availability-checking -parse-as-library)
+
+// This test uses `leaks` which is only available on apple platforms; limit it to macOS:
+// REQUIRES: OS=macosx
+
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+// REQUIRES: concurrency_runtime
+
+// UNSUPPORTED: back_deployment_runtime
+// UNSUPPORTED: freestanding
+
+// FIXME: enable discarding taskgroup on windows; rdar://104762037
+// UNSUPPORTED: OS=windows-msvc
+
+import _Concurrency
+
+struct Boom: Error {
+  let id: String
+
+  init(file: String = #fileID, line: UInt = #line) {
+    self.id = "\(file):\(line)"
+  }
+
+  init(id: String) {
+    self.id = id
+  }
+}
+
+struct IgnoredBoom: Error {}
+
+@discardableResult
+func echo(_ i: Int) -> Int { i }
+@discardableResult
+func boom(file: String = #fileID, line: UInt = #line) throws -> Int { throw Boom(file: file, line: line) }
+
+func shouldEqual<T: Equatable>(_ lhs: T, _ rhs: T) {
+  precondition(lhs == rhs, "'\(lhs)' did not equal '\(rhs)'")
+}
+func shouldStartWith(_ lhs: Any, _ rhs: Any) {
+  let l = "\(lhs)"
+  let r = "\(rhs)"
+  precondition(l.prefix("\(r)".count) == r, "'\(l)' did not start with '\(r)'")
+}
+
+final class SomeClass: @unchecked Sendable {
+  init() {}
+}
+
+// NOTE: Not as StdlibUnittest/TestSuite since these types of tests are unreasonably slow to load/debug.
+
+@main struct Main {
+  static func main() async {
+    // one error, one ok
+//    _ = try? await withThrowingDiscardingTaskGroup() { group in
+//      group.addTask {
+//        try boom()
+//      }
+//      group.addTask {
+//        return SomeClass() // will be discarded
+//      }
+//
+//      return 12
+//    }
+//
+//    // many ok
+//    _ = try? await withThrowingDiscardingTaskGroup() { group in
+//      for i in 0..<100 {
+//        group.addTask {
+//          return SomeClass() // will be discarded
+//        }
+//      }
+//
+//      return 12
+//    }
+
+    // many errors
+    _ = try? await withThrowingDiscardingTaskGroup() { group in
+      group.addTask {
+        try boom()
+      }
+      group.addTask {
+        try boom()
+      }
+      group.addTask {
+        try boom()
+      }
+      group.addTask {
+        try boom()
+      }
+
+      return 12
+    }
+//
+//    // many errors, many values
+//    _ = try? await withThrowingDiscardingTaskGroup() { group in
+//      group.addTask {
+//        return SomeClass() // will be discarded
+//      }
+//      group.addTask {
+//        return SomeClass() // will be discarded
+//      }
+//      group.addTask {
+//        return SomeClass() // will be discarded
+//      }
+//      group.addTask {
+//        try boom()
+//      }
+//      group.addTask {
+//        try boom()
+//      }
+//      group.addTask {
+//        try boom()
+//      }
+//      group.addTask {
+//        try boom()
+//      }
+//
+//      return 12
+//    }
+  }
+}