Fix waitUntil, which was broken in the new implementation of polling expectations

Sources/Nimble/Utils/AsyncAwait.swift

@@ -44,7 +44,7 @@ internal enum AsyncPollResult<T> {
     }
 }
 
-final class BlockingTask: Sendable {
+private final class BlockingTask: Sendable {
     private nonisolated(unsafe) var finished = false
     private nonisolated(unsafe) var continuation: CheckedContinuation<Void, Never>? = nil
     let sourceLocation: SourceLocation

Sources/Nimble/Utils/PollAwait.swift

@@ -81,6 +81,30 @@ internal enum PollStatus {
     case incomplete
 }
 
+private final class LockBox<T>: Sendable {
+    private nonisolated(unsafe) var value: T
+
+    private let lock = NSLock()
+
+    init(value: T) {
+        self.value = value
+    }
+
+    var currentValue: T {
+        lock.lock()
+        defer {
+            lock.unlock()
+        }
+        return value
+    }
+
+    func set(_ newValue: T) {
+        lock.lock()
+        value = newValue
+        lock.unlock()
+    }
+}
+
 func synchronousWaitUntil(
     timeout: NimbleTimeInterval,
     fnName: String,
@@ -100,23 +124,20 @@ Please use Swift Testing's `confirmation(...)` APIs to accomplish (nearly) the s
     return guaranteeNotNested(fnName: fnName, sourceLocation: sourceLocation) {
         let runloop = RunLoop.current
 
-        nonisolated(unsafe) var result = PollResult<Void>.timedOut
-        let lock = NSLock()
+        let resultBox = LockBox(value: PollResult<Void>.timedOut)
 
         let doneBlock: () -> Void = {
             let onFinish = {
-                lock.lock()
-                defer { lock.unlock() }
-                if case .completed = result {
+                if case .completed = resultBox.currentValue {
                     fail("waitUntil(...) expects its completion closure to be only called once", location: sourceLocation)
                     return
                 }
+                resultBox.set(.completed(()))
 #if canImport(CoreFoundation)
                 CFRunLoopStop(CFRunLoopGetCurrent())
 #else
                 RunLoop.main._stop()
 #endif
-                result = .completed(())
             }
             if Thread.isMainThread {
                 onFinish()
@@ -127,33 +148,30 @@ Please use Swift Testing's `confirmation(...)` APIs to accomplish (nearly) the s
 
         let capture = NMBExceptionCapture(
             handler: ({ exception in
-                lock.lock()
-                defer { lock.unlock() }
-                result = .raisedException(exception)
+                resultBox.set(.raisedException(exception))
             }),
             finally: ({ })
         )
         capture.tryBlock {
             do {
                 try closure(doneBlock)
             } catch {
-                lock.lock()
-                defer { lock.unlock() }
-                result = .errorThrown(error)
+                resultBox.set(.errorThrown(error))
             }
         }
 
-        if Thread.isMainThread {
-            runloop.run(mode: .default, before: Date(timeIntervalSinceNow: timeout.timeInterval))
-        } else {
-            DispatchQueue.main.sync {
-                _ = runloop.run(mode: .default, before: Date(timeIntervalSinceNow: timeout.timeInterval))
+        let start = Date()
+        while case .timedOut = resultBox.currentValue, abs(start.timeIntervalSinceNow) < timeout.timeInterval {
+            if Thread.isMainThread {
+                runloop.run(mode: .default, before: Date(timeIntervalSinceNow: timeout.timeInterval))
+            } else {
+                DispatchQueue.main.sync {
+                    _ = runloop.run(mode: .default, before: Date(timeIntervalSinceNow: timeout.timeInterval))
+                }
             }
         }
 
-        lock.lock()
-        defer { lock.unlock() }
-        return result
+        return resultBox.currentValue
     }
 }
 

Tests/NimbleTests/AsyncAwaitTest.swift

@@ -298,6 +298,74 @@ final class AsyncAwaitTest: XCTestCase { // swiftlint:disable:this type_body_len
         timer.cancel()
     }
 
+    func testWaitUntilNestedMainQueueAsyncCalls() async {
+        await waitUntil(timeout: .seconds(1)) { done in
+            DispatchQueue.main.async {
+                DispatchQueue.main.async {
+                    DispatchQueue.main.async {
+                        done()
+                    }
+                }
+            }
+        }
+    }
+
+    func testWaitUntilOperationQueueWithMainUnderlyingQueueAndBarrier() async {
+        await waitUntil(timeout: .seconds(1)) { done in
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = .main
+
+            let operation = BlockOperation {}
+
+            operationQueue.addOperation(operation)
+            operationQueue.addBarrierBlock {
+                done()
+            }
+        }
+    }
+
+    func testNestedMainAsyncWithOperationQueue() async {
+        await waitUntil(timeout: .seconds(1)) { done in
+            let indexingQueue = DispatchQueue.main
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = indexingQueue
+
+            indexingQueue.async {
+                indexingQueue.async {
+                    indexingQueue.async {
+                        let operation = BlockOperation {}
+
+                        operationQueue.addOperation(operation)
+                        operationQueue.addBarrierBlock {
+                            done()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    func testNestedBackgroundAsyncWithOperationQueue() async {
+        await waitUntil(timeout: .seconds(1)) { done in
+            let indexingQueue = DispatchQueue(label: "test.queue")
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = indexingQueue
+
+            indexingQueue.async {
+                indexingQueue.async {
+                    indexingQueue.async {
+                        let operation = BlockOperation {}
+
+                        operationQueue.addOperation(operation)
+                        operationQueue.addBarrierBlock {
+                            done()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     final class ClassUnderTest {
         var deinitCalled: (() -> Void)?
         var count = 0

Tests/NimbleTests/PollingTest.swift

@@ -129,7 +129,7 @@
             }
         }
     }
 
     func testToNeverDoesNotFailStalledMainThreadActivity() {
         func spinAndReturnTrue() -> Bool {
             Thread.sleep(forTimeInterval: 0.1)
@@ -137,7 +137,7 @@
         }
         expect(spinAndReturnTrue()).toNever(beFalse())
     }
 
     func testToAlwaysDetectsStalledMainThreadActivity() {
         func spinAndReturnTrue() -> Bool {
             Thread.sleep(forTimeInterval: 0.1)
@@ -208,7 +208,7 @@
 
         for index in 0..<1000 {
             if failed { break }
             waitUntil() { done in
                 DispatchQueue(label: "Nimble.waitUntilTest.\(index)").async {
                     done()
                 }
@@ -219,6 +219,74 @@
 #endif // canImport(Darwin)
     }
 
+    func testWaitUntilNestedMainQueueAsyncCalls() {
+        waitUntil(timeout: .seconds(1)) { done in
+            DispatchQueue.main.async {
+                DispatchQueue.main.async {
+                    DispatchQueue.main.async {
+                        done()
+                    }
+                }
+            }
+        }
+    }
+
+    func testWaitUntilOperationQueueWithMainUnderlyingQueueAndBarrier() {
+        waitUntil(timeout: .seconds(1)) { done in
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = .main
+
+            let operation = BlockOperation {}
+
+            operationQueue.addOperation(operation)
+            operationQueue.addBarrierBlock {
+                done()
+            }
+        }
+    }
+
+    func testNestedMainAsyncWithOperationQueue() {
+        waitUntil(timeout: .seconds(1)) { done in
+            let indexingQueue = DispatchQueue.main
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = indexingQueue
+
+            indexingQueue.async {
+                indexingQueue.async {
+                    indexingQueue.async {
+                        let operation = BlockOperation {}
+
+                        operationQueue.addOperation(operation)
+                        operationQueue.addBarrierBlock {
+                            done()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    func testNestedBackgroundAsyncWithOperationQueue() {
+        waitUntil(timeout: .seconds(1)) { done in
+            let indexingQueue = DispatchQueue(label: "test.queue")
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = indexingQueue
+
+            indexingQueue.async {
+                indexingQueue.async {
+                    indexingQueue.async {
+                        let operation = BlockOperation {}
+
+                        operationQueue.addOperation(operation)
+                        operationQueue.addBarrierBlock {
+                            done()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     func testWaitUntilAllowsInBackgroundThread() {
 #if !SWIFT_PACKAGE
         var executedAsyncBlock: Bool = false
@@ -343,4 +411,4 @@
     }
 }
 
 #endif // #if !os(WASI)