feat(minor): Use Swift 6's isolated(any) feature to unblock the main actor and run benchmarks in whatever isolation context they require (#358)

Sources/Benchmark/Benchmark.swift

@@ -8,7 +8,6 @@
 // http://www.apache.org/licenses/LICENSE-2.0
 //
 
-import Dispatch
 import Foundation
 
 // swiftlint:disable file_length identifier_name
@@ -18,11 +17,11 @@ public final class Benchmark: Codable, Hashable { // swiftlint:disable:this type
     @_documentation(visibility: internal)
     public typealias BenchmarkClosure = (_ benchmark: Benchmark) -> Void
     @_documentation(visibility: internal)
-    public typealias BenchmarkAsyncClosure = (_ benchmark: Benchmark) async -> Void
+    public typealias BenchmarkAsyncClosure = @isolated(any) (_ benchmark: Benchmark) async -> Void
     @_documentation(visibility: internal)
     public typealias BenchmarkThrowingClosure = (_ benchmark: Benchmark) throws -> Void
     @_documentation(visibility: internal)
-    public typealias BenchmarkAsyncThrowingClosure = (_ benchmark: Benchmark) async throws -> Void
+    public typealias BenchmarkAsyncThrowingClosure = @isolated(any) (_ benchmark: Benchmark) async throws -> Void
     @_documentation(visibility: internal)
     public typealias BenchmarkMeasurementSynchronization = (_ explicitStartStop: Bool) -> Void
     @_documentation(visibility: internal)
@@ -151,8 +150,11 @@ public final class Benchmark: Codable, Hashable { // swiftlint:disable:this type
     public var executablePath: String?
     /// closure: The actual benchmark closure that will be measured
     var closure: BenchmarkClosure? // The actual benchmark to run
-    /// asyncClosure: The actual benchmark (async) closure that will be measured
-    var asyncClosure: BenchmarkAsyncClosure? // The actual benchmark to run
+    /// asyncClosure: The actual benchmark (async) closure that will be measured.
+    /// Stored as the throwing variant so both throwing and non-throwing async closures can be held
+    /// while preserving the closure's actor isolation (captured via `@isolated(any)`), which lets the
+    /// executor run the measurement loop *on* that isolation hop-free (e.g. `@MainActor` benchmarks).
+    var asyncClosure: BenchmarkAsyncThrowingClosure? // The actual benchmark to run
     // setup/teardown hooks for the instance
     var setup: BenchmarkSetupHook?
     var teardown: BenchmarkTeardownHook?
@@ -324,26 +326,26 @@ public final class Benchmark: Codable, Hashable { // swiftlint:disable:this type
     ///   - setup: A closure that will be run once before the benchmark iterations are run
     ///   - teardown: A closure that will be run once after the benchmark iterations are done
     @discardableResult
-    public convenience init?(
+    public init?(
         _ name: String,
         configuration: Benchmark.Configuration = Benchmark.defaultConfiguration,
         closure: @escaping BenchmarkAsyncThrowingClosure,
         setup: BenchmarkSetupHook? = nil,
         teardown: BenchmarkTeardownHook? = nil
     ) {
-        self.init(
-            name,
-            configuration: configuration,
-            closure: { benchmark in
-                do {
-                    try await closure(benchmark)
-                } catch {
-                    benchmark.error("Benchmark \(name) failed with \(String(reflecting: error))")
-                }
-            },
-            setup: setup,
-            teardown: teardown
-        )
+        if configuration.skip {
+            return nil
+        }
+        target = ""
+        self.baseName = name
+        self.configuration = configuration
+        // Stored directly (rather than wrapped in an error-handling closure) so the closure's actor
+        // isolation is preserved; any thrown error is caught in `run()`.
+        asyncClosure = closure
+        self.setup = setup
+        self.teardown = teardown
+
+        benchmarkRegistration()
     }
 
     // Shared between sync/async actual benchmark registration
@@ -425,38 +427,27 @@ public final class Benchmark: Codable, Hashable { // swiftlint:disable:this type
     // The rest is intenral supporting infrastructure that should only
     // be used by the BenchmarkRunner
 
-    // https://forums.swift.org/t/actually-waiting-for-a-task/56230
-    // Async closures can possibly show false memory leaks possibly due to Swift runtime allocations
-    func runAsync() {
-        guard let asyncClosure else {
-            fatalError("Tried to runAsync on benchmark instance without any async closure set")
-        }
-
-        let semaphore = DispatchSemaphore(value: 0)
-
-        // Must do this in a separate thread, otherwise we block the concurrent thread pool
-        DispatchQueue.global(qos: .userInitiated)
-            .async {
-                Task {
-                    self._startMeasurement(false)
-                    await asyncClosure(self)
-                    self._stopMeasurement(false)
-
-                    semaphore.signal()
-                }
-            }
-        semaphore.wait()
-    }
-
     // Public but should only be used by BenchmarkRunner
+    //
+    // The `isolation` parameter defaults to the caller's isolation (`#isolation`), so when the
+    // executor runs the measurement loop on a benchmark's desired actor (e.g. `@MainActor`), the
+    // async closure below is invoked on that same isolation with no actor hop. There is no longer
+    // any thread-blocking bridge (the old `DispatchSemaphore`), so benchmarks may freely run on the
+    // main actor / main thread without deadlocking.
     @_documentation(visibility: internal)
-    public func run() {
+    public func run(isolation: isolated (any Actor)? = #isolation) async {
         if let closure {
             _startMeasurement(false)
             closure(self)
             _stopMeasurement(false)
-        } else {
-            runAsync()
+        } else if let asyncClosure {
+            _startMeasurement(false)
+            do {
+                try await asyncClosure(self)
+            } catch {
+                self.error("Benchmark \(name) failed with \(String(reflecting: error))")
+            }
+            _stopMeasurement(false)
         }
     }
 }

Sources/Benchmark/BenchmarkExecutor.swift

@@ -23,7 +23,13 @@ struct BenchmarkExecutor { // swiftlint:disable:this type_body_length
     let operatingSystemStatsProducer = OperatingSystemStatsProducer()
 
     // swiftlint:disable cyclomatic_complexity function_body_length
-    func run(_ benchmark: Benchmark) -> [BenchmarkResult] {
+    //
+    // The measurement loop runs isolated to the benchmark's desired actor (passed in by the
+    // caller from the closure's own `@isolated(any)` isolation).  For a `@MainActor` benchmark this
+    // means the entire loop — and each `await benchmark.run()` below — executes on the main actor
+    // with no per-iteration actor hop; for an unisolated benchmark `isolation` is `nil` and the loop
+    // runs on the generic executor.
+    func run(_ benchmark: Benchmark, isolation: isolated (any Actor)? = #isolation) async -> [BenchmarkResult] {
         var wallClockDuration: Duration = .zero
         var startMallocStats = MallocStats()
         var stopMallocStats = MallocStats()
@@ -52,7 +58,7 @@ struct BenchmarkExecutor { // swiftlint:disable:this type_body_length
 
         for iterations in 0..<benchmark.configuration.warmupIterations {
             benchmark.currentIteration = iterations
-            benchmark.run()
+            await benchmark.run()
         }
 
         #if canImport(OSLog)
@@ -395,7 +401,7 @@ struct BenchmarkExecutor { // swiftlint:disable:this type_body_length
 
             benchmark.currentIteration = iterations + benchmark.configuration.warmupIterations
 
-            benchmark.run()
+            await benchmark.run()
 
             iterations += 1
 

Sources/Benchmark/BenchmarkRunner.swift

@@ -217,7 +217,22 @@ public struct BenchmarkRunner: AsyncParsableCommand, BenchmarkRunnerReadWrite {
                             try suppressor.suppressOutput()
                         }
 
-                        results = benchmarkExecutor.run(benchmark)
+                        // Run the measurement loop on the benchmark's desired actor isolation, if any.
+                        // An `@isolated(any)` async closure carries its isolation (e.g. `@MainActor`);
+                        // unisolated and synchronous benchmarks yield `nil` and run on the generic executor.
+                        //
+                        // (Note: read `.isolation` off an unwrapped closure rather than via optional chaining,
+                        //        as the latter currently crashes the Swift compiler -
+                        //        https://github.com/swiftlang/swift/issues/89584).
+                        let isolation: (any Actor)?
+
+                        if let asyncClosure = benchmark.asyncClosure {
+                            isolation = asyncClosure.isolation
+                        } else {
+                            isolation = nil
+                        }
+
+                        results = await benchmarkExecutor.run(benchmark, isolation: isolation)
 
                         if quiet {
                             try suppressor.restoreOutput()

Tests/BenchmarkTests/BenchmarkTests.swift

@@ -8,28 +8,60 @@
 // http://www.apache.org/licenses/LICENSE-2.0
 //
 
+import Dispatch
 import XCTest
 
 @testable import Benchmark
 
 final class BenchmarkTests: XCTestCase {
-    func testBenchmarkRun() throws {
+    func testBenchmarkRun() async throws {
         let benchmark = Benchmark("testBenchmarkRun benchmark") { _ in
         }
         XCTAssertNotNil(benchmark)
-        benchmark?.run()
+        await benchmark?.run()
     }
 
-    func testBenchmarkRunAsync() throws {
+    func testBenchmarkRunAsync() async throws {
         func asyncFunc() async {}
         let benchmark = Benchmark("testBenchmarkRunAsync benchmark") { _ in
             await asyncFunc()
         }
         XCTAssertNotNil(benchmark)
-        benchmark?.runAsync()
+        await benchmark?.run()
     }
 
-    func testBenchmarkRunCustomMetric() throws {
+    func testBenchmarkRunAsyncThrowingFailure() async throws {
+        // Error handling for async throwing closures moved from an init-time wrapper into run();
+        // a thrown error must still be surfaced as a failureReason.
+        struct Boom: Error {}
+        let benchmark = Benchmark("testBenchmarkRunAsyncThrowingFailure benchmark") { _ in
+            await Task.yield()
+            throw Boom()
+        }
+        XCTAssertNotNil(benchmark)
+        await benchmark?.run()
+        XCTAssertNotNil(benchmark?.failureReason)
+    }
+
+    func testBenchmarkRunMainActor() async throws {
+        // An async benchmark whose closure is isolated to the main actor must carry that isolation,
+        // run on the main actor, and do so without deadlocking (the old semaphore bridge deadlocked).
+        let benchmark = Benchmark("testBenchmarkRunMainActor benchmark") { @MainActor _ in
+            await Task.yield()
+            MainActor.assertIsolated()
+        }
+        XCTAssertNotNil(benchmark)
+        let isolation: (any Actor)?
+        if let asyncClosure = benchmark?.asyncClosure {
+            isolation = asyncClosure.isolation
+        } else {
+            isolation = nil
+        }
+        XCTAssertIdentical(isolation, MainActor.shared)
+        await benchmark?.run(isolation: isolation)
+    }
+
+    func testBenchmarkRunCustomMetric() async throws {
         let benchmark = Benchmark(
             "testBenchmarkRunCustomMetric benchmark",
             configuration: .init(metrics: [.custom("customMetric")])
@@ -39,7 +71,7 @@ final class BenchmarkTests: XCTestCase {
             }
         }
         XCTAssertNotNil(benchmark)
-        benchmark?.run()
+        await benchmark?.run()
     }
 
     func testBenchmarkEqualityAndDifference() throws {
@@ -50,20 +82,20 @@ final class BenchmarkTests: XCTestCase {
         XCTAssertNotEqual(benchmark, benchmark2)
     }
 
-    func testBenchmarkRunFailure() throws {
+    func testBenchmarkRunFailure() async throws {
         let benchmark = Benchmark(
             "testBenchmarkRunFailure benchmark",
             configuration: .init(metrics: [.custom("customMetric")])
         ) { benchmark in
             benchmark.error("Benchmark failed")
         }
         XCTAssertNotNil(benchmark)
-        benchmark?.run()
+        await benchmark?.run()
         XCTAssertNotNil(benchmark?.failureReason)
         XCTAssertEqual(benchmark?.failureReason, "Benchmark failed")
     }
 
-    func testBenchmarkRunMoreParameters() throws {
+    func testBenchmarkRunMoreParameters() async throws {
         let benchmark = Benchmark(
             "testBenchmarkRunMoreParameters benchmark",
             configuration: .init(
@@ -78,7 +110,7 @@ final class BenchmarkTests: XCTestCase {
             }
         }
         XCTAssertNotNil(benchmark)
-        benchmark?.run()
+        await benchmark?.run()
     }
 
     func testBenchmarkParameterizedDescription() throws {
@@ -95,4 +127,97 @@ final class BenchmarkTests: XCTestCase {
         XCTAssertNotNil(benchmark)
         XCTAssertEqual(benchmark?.name, "testBenchmarkParameterizedDescription benchmark (bin: 42, foo: bar, pi: 3.14)")
     }
+
+    func testExecutorRunsMainActorBenchmarkWithoutDeadlock() async throws {
+        // End-to-end through the real measurement loop: a main-actor-isolated async benchmark must
+        // complete (the old DispatchSemaphore bridge deadlocked here) and actually run on the main actor.
+        Benchmark.testSkipBenchmarkRegistrations = true
+        defer { Benchmark.testSkipBenchmarkRegistrations = false }
+
+        final class Counter: @unchecked Sendable { var count = 0 }
+        let counter = Counter()
+
+        let benchmark = Benchmark(
+            "MainActor executor benchmark",
+            configuration: .init(metrics: [.wallClock], warmupIterations: 1, maxDuration: .seconds(10), maxIterations: 3)
+        ) { @MainActor _ in
+            await Task.yield()
+            MainActor.assertIsolated()
+            counter.count += 1
+        }
+        XCTAssertNotNil(benchmark)
+
+        let results = await BenchmarkExecutor().run(benchmark!, isolation: MainActor.shared)
+        XCTAssertFalse(results.isEmpty)
+        XCTAssertGreaterThan(counter.count, 0)
+    }
+
+    func testUnisolatedAsyncBenchmarkHoppingToMainActor() async throws {
+        // A *nominally unisolated* async benchmark that hops to the main actor itself (via MainActor.run)
+        // must complete without deadlocking.  This is one of the scenarios the old DispatchSemaphore
+        // bridge could deadlock on.  The closure carries no isolation, so the executor runs unisolated
+        // (isolation == nil) — exactly as BenchmarkRunner drives it — and the hop happens inside.
+        Benchmark.testSkipBenchmarkRegistrations = true
+        defer { Benchmark.testSkipBenchmarkRegistrations = false }
+
+        final class Flag: @unchecked Sendable { var ran = false }
+        let flag = Flag()
+
+        let benchmark = Benchmark(
+            "unisolated async hops to MainActor",
+            configuration: .init(metrics: [.wallClock], warmupIterations: 1, maxDuration: .seconds(10), maxIterations: 3)
+        ) { _ in
+            await MainActor.run {
+                MainActor.assertIsolated()
+                flag.ran = true
+            }
+        }
+        XCTAssertNotNil(benchmark)
+
+        // An unisolated async closure carries no isolation.
+        let isolation: (any Actor)?
+        if let asyncClosure = benchmark?.asyncClosure {
+            isolation = asyncClosure.isolation
+        } else {
+            isolation = nil
+        }
+        XCTAssertNil(isolation)
+
+        let results = await BenchmarkExecutor().run(benchmark!, isolation: isolation)
+        XCTAssertFalse(results.isEmpty)
+        XCTAssertTrue(flag.ran)
+    }
+
+    func testUnisolatedSyncBenchmarkCallingMainSync() {
+        // A *synchronous* benchmark that bounces to the main queue (DispatchQueue.main.sync) must run
+        // off the main thread and not deadlock.  Driven from a detached task (off-main) while the main
+        // thread drains the main queue in wait(for:).  An expectation+timeout means that, should this
+        // ever regress to a deadlock, the test eventually fails outright rather than hanging the whole
+        // suite.
+        Benchmark.testSkipBenchmarkRegistrations = true
+        defer { Benchmark.testSkipBenchmarkRegistrations = false }
+
+        final class Flag: @unchecked Sendable { var ran = false }
+        let flag = Flag()
+
+        let benchmark = Benchmark(
+            "unisolated sync calls DispatchQueue.main.sync",
+            configuration: .init(metrics: [.wallClock], warmupIterations: 1, maxDuration: .seconds(10), maxIterations: 3)
+        ) { _ in
+            DispatchQueue.main.sync {
+                // NB: we only record that the block ran, not the thread identity — on Linux the main
+                // queue is drained off the main thread, so Thread.isMainThread would not be reliable.
+                flag.ran = true
+            }
+        }
+        XCTAssertNotNil(benchmark)
+
+        let completed = expectation(description: "benchmark completed without deadlocking")
+        Task.detached {
+            _ = await BenchmarkExecutor().run(benchmark!, isolation: nil)
+            completed.fulfill()
+        }
+        wait(for: [completed], timeout: 30)
+        XCTAssertTrue(flag.ran)
+    }
 }