From b814a514c2e50438322473ca32e63678225bf677 Mon Sep 17 00:00:00 2001
From: Konrad `ktoso` Malawski <konrad_malawski@apple.com>
Date: Wed, 19 Jul 2023 18:28:37 +0900
Subject: [PATCH 1/2] Test reenable:
 async_taskgroup_discarding_dontLeak_class_error.swift

---
 ...roup_discarding_dontLeak_class_error.swift | 56 ++++++++++++++++---
 1 file changed, 48 insertions(+), 8 deletions(-)

diff --git a/test/Concurrency/Runtime/async_taskgroup_discarding_dontLeak_class_error.swift b/test/Concurrency/Runtime/async_taskgroup_discarding_dontLeak_class_error.swift
index 8227662a7b375..b22c55bf5ced7 100644
--- a/test/Concurrency/Runtime/async_taskgroup_discarding_dontLeak_class_error.swift
+++ b/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
   }
-}
+}
\ No newline at end of file

From d01abcb42e9ad7f407b394a8e03b869d7e54fc68 Mon Sep 17 00:00:00 2001
From: Konrad `ktoso` Malawski <konrad_malawski@apple.com>
Date: Thu, 3 Aug 2023 16:59:36 +0900
Subject: [PATCH 2/2] Unlock BEFORE resuming the waitAll waiting task, to avoid
 unlock() on released memory.

Because the waitAll() group method is called before the group returns
from withTaskGroup and is used to ensure all tasks have completed before
we destroy the task group, this method is then immediately followed by
calling TaskGroup::destroy.

If we, as previously, first resume the waiting task and then attempt to
unlock the lock held by the group, we are in an unsafe situation with
racing the task group destroy() and the unlock().

Therefore, we must release the lock before we resume the waiting task.
---
 stdlib/public/Concurrency/TaskGroup.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/stdlib/public/Concurrency/TaskGroup.cpp b/stdlib/public/Concurrency/TaskGroup.cpp
index b17dd31a5a886..03f726b65b05e 100644
--- a/stdlib/public/Concurrency/TaskGroup.cpp
+++ b/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;
   }