Merge pull request #67417 from ktoso/wip-reenable-async_taskgroup_discarding

Author: ktoso

test/Concurrency/Runtime/async_taskgroup_discarding.swift

@@ -16,29 +16,75 @@ import _Concurrency
 final class PayloadFirst {}
 final class PayloadSecond {}
 
+actor SimpleCountDownLatch {
+  let from: Int
+  var count: Int
+
+  var continuation: CheckedContinuation<Void, Never>?
+
+  init(from: Int) {
+    self.from = from
+    self.count = from
+  }
+
+  func hit() {
+    defer { count -= 1 }
+    print("hit @ \(count)")
+    if count == 0 {
+      fatalError("Counted down more times than expected! (From: \(from))")
+    } else if count == 1 {
+      print("hit resume")
+      continuation?.resume()
+    }
+  }
+
+  func wait() async {
+    guard self.count > 0 else {
+      return // we're done
+    }
+
+    return await withCheckedContinuation { cc in
+      self.continuation = cc
+    }
+  }
+}
+
 final class ErrorFirst: Error {
   let first: PayloadFirst
+  let id: String
+  let latch: SimpleCountDownLatch
 
-  init(file: String = #fileID, line: UInt = #line) {
+  init(latch: SimpleCountDownLatch, file: String = #fileID, line: UInt = #line) {
+    self.latch = latch
+    self.id = "\(file):\(line)"
     first = .init()
+    print("init \(self) id:\(id)")
   }
   deinit {
-    print("deinit \(self)")
+    print("deinit \(self) id:\(id)")
+    Task { [latch] in await latch.hit() }
   }
 }
 
+// Should not really matter that different types, but want to make really sure
 final class ErrorSecond: Error {
-  let second: PayloadSecond
+  let first: PayloadFirst
+  let id: String
+  let latch: SimpleCountDownLatch
 
-  init(file: String = #fileID, line: UInt = #line) {
-    second = .init()
+  init(latch: SimpleCountDownLatch, file: String = #fileID, line: UInt = #line) {
+    self.latch = latch
+    self.id = "\(file):\(line)"
+    first = .init()
+    print("init \(self) id:\(id)")
   }
-
   deinit {
-    print("deinit \(self)")
+    print("deinit \(self) id:\(id)")
+    Task { [latch] in await latch.hit() }
   }
 }
 
+
 func shouldStartWith(_ lhs: Any, _ rhs: Any) {
   let l = "\(lhs)"
   let r = "\(rhs)"
@@ -47,20 +93,25 @@ func shouldStartWith(_ lhs: Any, _ rhs: Any) {
 
 // NOTE: Not as StdlibUnittest/TestSuite since these types of tests are unreasonably slow to load/debug.
 
+@discardableResult
+func one() -> Int {
+  1
+}
+
 @main struct Main {
   static func main() async {
+    let latch = SimpleCountDownLatch(from: 6)
     do {
+
       let got = try await withThrowingDiscardingTaskGroup() { group in
-        group.addTask {
-          1
-        }
-        group.addTask {
-          throw ErrorFirst()
-        }
+        group.addTask { one() }
+        group.addTask { throw ErrorFirst(latch: latch) }
+        group.addTask { throw ErrorFirst(latch: latch) }
+        group.addTask { throw ErrorFirst(latch: latch) }
 
-        group.addTask {
-          throw ErrorSecond()
-        }
+        group.addTask { throw ErrorSecond(latch: latch) }
+        group.addTask { throw ErrorSecond(latch: latch) }
+        group.addTask { throw ErrorSecond(latch: latch) }
 
         return 12
       }
@@ -70,5 +121,12 @@ func shouldStartWith(_ lhs: Any, _ rhs: Any) {
     }
     // CHECK: deinit main.Error
     // CHECK: deinit main.Error
+    // CHECK: deinit main.Error
+
+    // CHECK: deinit main.Error
+    // CHECK: deinit main.Error
+    // CHECK: deinit main.Error
+    await latch.wait()
+    print("done") // CHECK: done
   }
 }