From 9a5cd297c6535986fe5b3e5b3198703fe02e677b Mon Sep 17 00:00:00 2001
From: dimlio <122263440+dimlio@users.noreply.github.com>
Date: Thu, 30 Mar 2023 22:55:58 +0300
Subject: [PATCH] chore: make startup and shutdown hooks async (and throwing)
 (#141)

Make benchmark startup and shutdown hook closures async and throwing.

---------

Co-authored-by: Joakim Hassila <jocke@ordo.one>
---
 .../BenchmarkTool/BenchmarkTool+Operations.swift  |  9 ++++-----
 Sources/Benchmark/Benchmark.swift                 |  2 +-
 Sources/Benchmark/BenchmarkRunner.swift           | 15 +++++++++++++--
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/Plugins/BenchmarkTool/BenchmarkTool+Operations.swift b/Plugins/BenchmarkTool/BenchmarkTool+Operations.swift
index 714eac73..0d8ac751 100644
--- a/Plugins/BenchmarkTool/BenchmarkTool+Operations.swift
+++ b/Plugins/BenchmarkTool/BenchmarkTool+Operations.swift
@@ -38,6 +38,9 @@ extension BenchmarkTool {
                 benchmarks.append(benchmark)
             case .end:
                 break outerloop
+            case let .error(description):
+                failBenchmark(description)
+                break outerloop
             default:
                 print("Unexpected reply \(benchmarkReply)")
             }
@@ -59,11 +62,7 @@ extension BenchmarkTool {
             case .end:
                 break outerloop
             case let .error(description):
-                print("*****")
-                print("***** Benchmark '\(benchmark.name)' failed:")
-                print("***** \(description)")
-                print("*****")
-                failBenchmark("")
+                failBenchmark(description)
                 break outerloop
             default:
                 print("Unexpected reply \(benchmarkReply)")
diff --git a/Sources/Benchmark/Benchmark.swift b/Sources/Benchmark/Benchmark.swift
index ee1327b4..cb436971 100644
--- a/Sources/Benchmark/Benchmark.swift
+++ b/Sources/Benchmark/Benchmark.swift
@@ -40,7 +40,7 @@ public final class Benchmark: Codable, Hashable {
     public typealias BenchmarkCustomMetricMeasurement = (BenchmarkMetric, Int) -> Void
 
     /// Alias for closures used to hook into setup / teardown
-    public typealias BenchmarkHook = () -> Void
+    public typealias BenchmarkHook = () async throws -> Void
 
     /// This closure if set, will be run before a targets benchmarks are run, but after they are registered
     public static var startupHook: BenchmarkHook?
diff --git a/Sources/Benchmark/BenchmarkRunner.swift b/Sources/Benchmark/BenchmarkRunner.swift
index de4497eb..ecbf8a49 100644
--- a/Sources/Benchmark/BenchmarkRunner.swift
+++ b/Sources/Benchmark/BenchmarkRunner.swift
@@ -113,7 +113,12 @@ public struct BenchmarkRunner: AsyncParsableCommand, BenchmarkRunnerReadWrite {
         var benchmark: Benchmark?
         var results: [BenchmarkResult] = []
 
-        Benchmark.startupHook?()
+        do {
+            try await Benchmark.startupHook?()
+        } catch {
+            try channel.write(.error("Benchmark.startupHook failed: \(error)"))
+            return
+        }
 
         while true {
             if debug { // in debug mode we run all benchmarks matching filter/skip specified
@@ -197,9 +202,15 @@ public struct BenchmarkRunner: AsyncParsableCommand, BenchmarkRunnerReadWrite {
                     print("Internal error: Couldn't find specified benchmark '\(benchmarkToRun.name)' to run.")
                 }
 
+                do {
+                    try await Benchmark.shutdownHook?()
+                } catch {
+                    try channel.write(.error("Benchmark.shutdownHook failed: \(error)"))
+                    return
+                }
+
                 try channel.write(.end)
             case .end:
-                Benchmark.shutdownHook?()
                 return
             }
         }