Merge pull request #67700 from ktoso/pick-reenable-async_taskgroup_discarding_dontLeak_class_error.swift

🍒[5.9][TaskGroup] Reenable test and fix memory issue

Author: DougGregor

stdlib/public/Concurrency/TaskGroup.cpp

@@ -1810,9 +1810,12 @@ void TaskGroupBase::waitAll(SwiftError* bodyError, AsyncTask *waitingTask,
       swift_release(completedTask);
     }
 
-    waitingTask->runInFullyEstablishedContext();
-
+    // We MUST release the lock before we resume the waiting task, because the resumption
+    // will allow it to destroy the task group, in which case the unlock()
+    // would be performed on freed memory (!)
     unlock();
+
+    waitingTask->runInFullyEstablishedContext();
     return;
   }
 

test/Concurrency/Runtime/async_taskgroup_discarding_dontLeak_class_error.swift

@@ -12,35 +12,75 @@
 
 import _Concurrency
 
+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 }
+    if count == 0 {
+      fatalError("Counted down more times than expected! (From: \(from))")
+    } else if count == 1 {
+      continuation?.resume()
+    }
+  }
+
+  func wait() async {
+    guard self.count > 0 else {
+      return // we're done
+    }
+
+    return await withCheckedContinuation { cc in
+      self.continuation = cc
+    }
+  }
+}
+
 final class ClassBoom: Error {
   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)"
     print("INIT OF ClassBoom from \(id)")
   }
 
   deinit {
     print("DEINIT OF ClassBoom from \(id)")
+    Task { [latch] in await latch.hit() }
   }
 }
 
 @main struct Main {
   static func main() async {
+    let latch = SimpleCountDownLatch(from: 4)
 
     // many errors
     _ = try? await withThrowingDiscardingTaskGroup() { group in
-      group.addTask { throw ClassBoom() }
-      group.addTask { throw ClassBoom() }
-      group.addTask { throw ClassBoom() }
-      group.addTask { throw ClassBoom() }
-      group.addTask { 12 }
-      return 12
+      group.addTask { throw ClassBoom(latch: latch) }
+      group.addTask { throw ClassBoom(latch: latch) }
+      group.addTask { throw ClassBoom(latch: latch) }
+      group.addTask { throw ClassBoom(latch: latch) }
+      group.addTask {
+        12 // ignore this on purpose
+      }
+      return 42
 
       // CHECK: DEINIT OF ClassBoom
       // CHECK: DEINIT OF ClassBoom
       // CHECK: DEINIT OF ClassBoom
       // CHECK: DEINIT OF ClassBoom
     }
+
+    await latch.wait()
+    print("done") // CHECK: done
   }
-}
+}