From 78f945647cb20503444b9ebd56ba38e5e28ae96c Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Fri, 11 Aug 2023 01:46:55 +0200
Subject: [PATCH] Add --local-runner cli flag

---
 src/args.rs   | 12 +++++++++---
 src/config.rs |  2 ++
 src/main.rs   | 30 +++++++++++++++++++++++++++++-
 src/runner.rs | 12 +++++++++++-
 src/server.rs |  6 +++++-
 5 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/src/args.rs b/src/args.rs
index 1e16422..6f515c4 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -5,24 +5,30 @@ pub const VERSION: &str = concat!(env!("CARGO_PKG_VERSION"), " (", env!("VERGEN_
 
 #[derive(Debug, clap::Parser)]
 pub struct ServerCommand {
-    /// Path to the repo's tablejohn database.
+    /// Path to a tablejohn database.
     pub db: PathBuf,
 
-    /// Path to the git repo.
+    /// Path to a git repo.
     pub repo: Option<PathBuf>,
 
-    /// Path to the bench repo.
+    /// Path to a bench repo.
     #[arg(long, short)]
     pub bench_repo: Option<PathBuf>,
 
     /// Open the UI in your browser.
     #[arg(long, short)]
     pub open: bool,
+
+    /// Start one or more local runners for this server.
+    #[arg(long, short, action = clap::ArgAction::Count)]
+    pub local_runner: u8,
 }
 
 #[derive(Debug, clap::Parser)]
 pub enum Command {
+    /// Start a tablejohn server.
     Server(ServerCommand),
+    /// Start a tablejohn benchmark runner.
     Runner,
     // TODO bench script command?
 }
diff --git a/src/config.rs b/src/config.rs
index 450085d..78a9af3 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -198,12 +198,14 @@ impl ConfigFile {
     }
 }
 
+#[derive(Clone)]
 pub struct RunnerServerConfig {
     /// Always ends with a `/`.
     pub url: String,
     pub token: String,
 }
 
+#[derive(Clone)]
 pub struct Config {
     /// Always starts and ends with a `/`.
     pub web_base: String,
diff --git a/src/main.rs b/src/main.rs
index 928475a..387460e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,6 +9,7 @@ mod somehow;
 use std::{io, process, time::Duration};
 
 use clap::Parser;
+use config::RunnerServerConfig;
 use tokio::{select, signal::unix::SignalKind};
 use tracing::{debug, error, info, Level};
 use tracing_subscriber::{
@@ -92,6 +93,30 @@ async fn open_in_browser(config: &Config) {
     }
 }
 
+async fn launch_local_runners(config: &'static Config, amount: u8) {
+    let server_name = "localhost";
+    let server_config = Box::leak(Box::new(RunnerServerConfig {
+        url: format!("http://{}{}", config.web_address, config.web_base),
+        token: config.web_runner_token.clone(),
+    }));
+
+    // Wait a bit to ensure the server is ready to serve requests.
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    for i in 0..amount {
+        let mut runner_config = config.clone();
+        runner_config.runner_name = format!("{}-{i}", config.runner_name);
+        let runner_config = Box::leak(Box::new(runner_config));
+
+        info!("Launching local runner {}", runner_config.runner_name);
+        runner::launch_standalone_server_task(
+            runner_config,
+            server_name.to_string(),
+            server_config,
+        );
+    }
+}
+
 async fn run() -> somehow::Result<()> {
     let args = Args::parse();
 
@@ -106,8 +131,11 @@ async fn run() -> somehow::Result<()> {
                 tokio::task::spawn(open_in_browser(config));
             }
 
-            let server = Server::new(config, command).await?;
+            if command.local_runner > 0 {
+                tokio::task::spawn(launch_local_runners(config, command.local_runner));
+            }
 
+            let server = Server::new(config, command).await?;
             select! {
                 _ = wait_for_signal() => {}
                 _ = server.run() => {}
diff --git a/src/runner.rs b/src/runner.rs
index d00ea54..189afe5 100644
--- a/src/runner.rs
+++ b/src/runner.rs
@@ -6,7 +6,7 @@ use std::sync::{Arc, Mutex};
 use tokio::task::JoinSet;
 use tracing::{debug, error};
 
-use crate::config::Config;
+use crate::config::{Config, RunnerServerConfig};
 
 use self::{coordinator::Coordinator, server::Server};
 
@@ -42,3 +42,13 @@ impl Runner {
         while tasks.join_next().await.is_some() {}
     }
 }
+
+pub fn launch_standalone_server_task(
+    config: &'static Config,
+    server_name: String,
+    server_config: &'static RunnerServerConfig,
+) {
+    let coordinator = Arc::new(Mutex::new(Coordinator::new()));
+    let mut server = Server::new(server_name, config, server_config, coordinator);
+    tokio::task::spawn(async move { server.run().await });
+}
diff --git a/src/server.rs b/src/server.rs
index c11328e..0588e8c 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -18,7 +18,11 @@ use sqlx::{
 use tokio::select;
 use tracing::{debug, info};
 
-use crate::{args::ServerCommand, config::Config, somehow};
+use crate::{
+    args::ServerCommand,
+    config::{Config, RunnerServerConfig},
+    runner, somehow,
+};
 
 use self::runners::Runners;