From 5bbe6da0b42c3c231d8298db8e4295f15401656a Mon Sep 17 00:00:00 2001
From: Thomas Leonard <talex5@gmail.com>
Date: Tue, 30 May 2023 11:30:29 +0100
Subject: [PATCH] Allow selecting a benchmark

---
 Makefile            |  2 +-
 bench/bench_main.ml | 25 -------------------------
 bench/dune          |  2 +-
 bench/main.ml       | 43 +++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+), 27 deletions(-)
 delete mode 100644 bench/bench_main.ml
 create mode 100644 bench/main.ml

diff --git a/Makefile b/Makefile
index 85b883d25..0a0d8084a 100644
--- a/Makefile
+++ b/Makefile
@@ -4,7 +4,7 @@ all:
 	dune build @runtest @all
 
 bench:
-	dune exec -- ./bench/bench_main.exe
+	dune exec -- ./bench/main.exe
 
 test_posix:
 	EIO_BACKEND=posix dune runtest
diff --git a/bench/bench_main.ml b/bench/bench_main.ml
deleted file mode 100644
index b7c8163be..000000000
--- a/bench/bench_main.ml
+++ /dev/null
@@ -1,25 +0,0 @@
-let benchmarks = [
-  "buf", Bench_buf_read.run;
-  "cancel", Bench_cancel.run;
-  "fd", Bench_fd.run;
-  "condition", Bench_condition.run;
-  "yield", Bench_yield.run;
-  "mutex", Bench_mutex.run;
-  "semaphore", Bench_semaphore.run;
-  "stream", Bench_stream.run;
-  "promise", Bench_promise.run;
-  "http", Bench_http.run;
-]
-
-let () =
-  Eio_main.run @@ fun env ->
-  let run (name, fn) =
-    let metrics = fn env in
-    `Assoc [
-      "name", `String name;
-      "metrics", `List (List.map Metric.to_json metrics);
-    ]
-  in
-  Yojson.Safe.pretty_to_channel stdout @@ `Assoc [
-    "results", `List (List.map run benchmarks);
-  ]
diff --git a/bench/dune b/bench/dune
index 4ab30361f..8f7902054 100644
--- a/bench/dune
+++ b/bench/dune
@@ -1,3 +1,3 @@
 (executables
-  (names bench_main)
+  (names main)
   (libraries eio_main yojson))
diff --git a/bench/main.ml b/bench/main.ml
new file mode 100644
index 000000000..6941b3c39
--- /dev/null
+++ b/bench/main.ml
@@ -0,0 +1,43 @@
+let benchmarks = [
+  "buf", Bench_buf_read.run;
+  "cancel", Bench_cancel.run;
+  "fd", Bench_fd.run;
+  "condition", Bench_condition.run;
+  "yield", Bench_yield.run;
+  "mutex", Bench_mutex.run;
+  "semaphore", Bench_semaphore.run;
+  "stream", Bench_stream.run;
+  "promise", Bench_promise.run;
+  "http", Bench_http.run;
+]
+
+let usage_error () =
+  let names = List.map fst benchmarks in
+  Fmt.epr "Usage: main.exe [%a]@." Fmt.(list ~sep:(any " | ") string) names;
+  exit 1
+
+let () =
+  Eio_main.run @@ fun env ->
+  let benchmarks =
+    match Array.to_list Sys.argv with
+    | [_] -> benchmarks
+    | [_; name] ->
+      begin match List.assoc_opt name benchmarks with
+        | Some run -> [name, run]
+        | None ->
+          Fmt.epr "Unknown benchmark %S@." name;
+          usage_error ()
+      end
+    | _ -> usage_error ()
+  in
+  let run (name, fn) =
+    Eio.traceln "Running %s..." name;
+    let metrics = fn env in
+    `Assoc [
+      "name", `String name;
+      "metrics", `List (List.map Metric.to_json metrics);
+    ]
+  in
+  Fmt.pr "%a@." (Yojson.Safe.pretty_print ~std:true) @@ `Assoc [
+    "results", `List (List.map run benchmarks);
+  ]