Fix async waitUntil waiting until whenever the closure finishes before returning

Author: younata

Sources/Nimble/Utils/AsyncAwait.swift

@@ -61,18 +61,30 @@ final class BlockingTask: Sendable {
     }
 
     func run() async {
-        if let continuation = lock.withLock({ self.continuation }) {
+        let continuation: CheckedContinuation<Void, Never>? = {
+            lock.lock()
+            let continuation = self.continuation
+            lock.unlock()
+            return continuation
+        }()
+
+        if let continuation {
             continuation.resume()
         }
         await withTaskCancellationHandler {
             await withCheckedContinuation {
                 lock.lock()
-                defer { lock.unlock() }
 
+                let shouldResume: Bool
                 if finished {
-                    $0.resume()
+                    shouldResume = true
                 } else {
                     self.continuation = $0
+                    shouldResume = false
+                }
+                lock.unlock()
+                if shouldResume {
+                    $0.resume()
                 }
             }
         } onCancel: {
@@ -83,25 +95,27 @@ final class BlockingTask: Sendable {
 
     func complete() {
         lock.lock()
-        defer { lock.unlock() }
+        let wasFinished = finished
+        finished = true
+        lock.unlock()
 
-        if finished {
+        if wasFinished {
             fail(
                 "waitUntil(...) expects its completion closure to be only called once",
                 location: sourceLocation
             )
         } else {
-            finished = true
             self.continuation?.resume()
             self.continuation = nil
         }
     }
 
     func handleCancellation() {
         lock.lock()
-        defer { lock.unlock() }
+        let wasFinished = finished
+        lock.unlock()
 
-        guard finished == false else {
+        guard wasFinished == false else {
             return
         }
         continuation?.resume()
@@ -156,15 +170,15 @@ internal func performBlock(
 #endif
     #endif
 
-    return await withTaskGroup(of: Void.self) { taskGroup in
+    return await withTaskGroup(of: Void.self, returning: AsyncPollResult<Void>.self) { taskGroup in
         let blocker = BlockingTask(sourceLocation: sourceLocation)
         let tracker = ResultTracker<Void>()
 
         taskGroup.addTask {
             await blocker.run()
         }
 
-        taskGroup.addTask {
+        let task = Task {
             do {
                 try await closure {
                     blocker.complete()
@@ -179,9 +193,7 @@ internal func performBlock(
             do {
                 try await Task.sleep(nanoseconds: (timeout + leeway).nanoseconds)
                 tracker.finish(with: .timedOut)
-            } catch {
-
-            }
+            } catch {}
         }
 
         var result: AsyncPollResult<Void> = .incomplete
@@ -194,6 +206,7 @@ internal func performBlock(
             break
         }
         taskGroup.cancelAll()
+        task.cancel()
         return result
     }
 }

Tests/NimbleTests/AsyncAwaitTest.swift

@@ -282,10 +282,14 @@ final class AsyncAwaitTest: XCTestCase { // swiftlint:disable:this type_body_len
         }
         timer.resume()
 
-        for index in 0..<1000 {
-            if failed { break }
+        let runQueue = DispatchQueue(label: "Nimble.waitUntilTest.runQueue", attributes: .concurrent)
+
+        for _ in 0..<1000 {
+            if failed {
+                break
+            }
             await waitUntil() { done in
-                DispatchQueue(label: "Nimble.waitUntilTest.\(index)").async {
+                runQueue.async {
                     done()
                 }
             }