taktwerk 0.8.0

//! Runs started from the workbench: each a `taktwerk run` child with an `http` connector on a
//! free loopback port and a token of its own, its API proxied under `/api/runs/<id>/live/`.

use std::collections::VecDeque;
use std::path::{Component, Path, PathBuf};
use std::process::Stdio;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use axum::body::Body;
use axum::http::{HeaderValue, Method, Request, StatusCode, header};
use axum::response::{IntoResponse, Response};
use hyper_util::rt::TokioIo;
use serde::Deserialize;
use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt as _, BufReader};
use tokio::process::Command;
use tokio::sync::oneshot;
use tracing::warn;

use super::files::{relative, resolve};
use super::unix_ms;

/// Log lines kept per run.
const LOG_LINES: usize = 400;
/// How long a stop request waits for the run to end before killing it.
const STOP_GRACE: Duration = Duration::from_secs(8);

/// What `POST /api/runs` asks for; paths relative to the folder.
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct StartRequest {
    /// Project file.
    pub project: String,
    /// Scenario file.
    #[serde(default)]
    pub scenario: Option<String>,
    /// `--speed`: `1`, a factor, or `max`.
    #[serde(default)]
    pub speed: Option<String>,
    /// `--cycles`.
    #[serde(default)]
    pub cycles: Option<u64>,
    /// `--until-time`, seconds.
    #[serde(default)]
    pub until_time: Option<f64>,
    /// `--set KEY=VALUE`, in order.
    #[serde(default)]
    pub set: Vec<String>,
    /// `--seed`.
    #[serde(default)]
    pub seed: Option<u64>,
    /// `--record-dir`, created when missing.
    #[serde(default)]
    pub record_dir: Option<String>,
    /// Start the project's own connectors; off by default, so a workbench run talks to no
    /// plant.
    #[serde(default)]
    pub connectors: bool,
}

/// The `taktwerk run` arguments for `req`, the folder's paths resolved; the `--http` ones are
/// added by the caller.
///
/// # Errors
/// A path outside the folder or missing, or a malformed option, as text.
pub fn arguments(root: &Path, req: &StartRequest) -> Result<Vec<String>, String> {
    let file = |rel: &str, what: &str| {
        resolve(root, rel)
            .filter(|p| p.is_file())
            .map(|p| p.display().to_string())
            .ok_or_else(|| format!("{what} `{rel}` is not a file in the folder"))
    };
    let mut args = vec!["run".to_owned(), file(&req.project, "project")?];
    if !req.connectors {
        args.push("--no-connectors".into());
    }
    if let Some(s) = &req.scenario {
        args.extend(["--scenario".into(), file(s, "scenario")?]);
    }
    if let Some(speed) = &req.speed {
        let ok = speed == "max" || speed.parse::<f64>().is_ok_and(|f| f > 0.0 && f.is_finite());
        if !ok {
            return Err(format!(
                "speed `{speed}`: expected `max` or a positive factor"
            ));
        }
        args.extend(["--speed".into(), speed.clone()]);
    }
    if let Some(n) = req.cycles {
        args.extend(["--cycles".into(), n.to_string()]);
    }
    if let Some(t) = req.until_time {
        if !t.is_finite() {
            return Err("until_time must be finite".into());
        }
        args.extend(["--until-time".into(), t.to_string()]);
    }
    for s in &req.set {
        if !s.contains('=') {
            return Err(format!("set `{s}`: expected KEY=VALUE"));
        }
        args.extend(["--set".into(), s.clone()]);
    }
    if let Some(seed) = req.seed {
        args.extend(["--seed".into(), seed.to_string()]);
    }
    if let Some(dir) = &req.record_dir {
        let rel = Path::new(dir);
        if rel.is_absolute() || rel.components().any(|c| !matches!(c, Component::Normal(_))) {
            return Err(format!(
                "record_dir `{dir}` must be a plain path inside the folder"
            ));
        }
        args.extend(["--record-dir".into(), root.join(rel).display().to_string()]);
    }
    Ok(args)
}

/// One run started from the workbench.
pub struct Launched {
    id: u32,
    project: String,
    args: Vec<String>,
    port: u16,
    token: String,
    started_ms: u64,
    state: Mutex<Ended>,
    log: Mutex<VecDeque<String>>,
    kill: Mutex<Option<oneshot::Sender<()>>>,
}

/// How a run ended, once it did.
#[derive(Debug, Clone, Default)]
struct Ended {
    at_ms: Option<u64>,
    code: Option<i32>,
}

impl Launched {
    fn running(&self) -> bool {
        self.state
            .lock()
            .map(|s| s.at_ms.is_none())
            .unwrap_or(false)
    }

    /// Summary as JSON; with the log tail when `log`.
    pub fn to_json(&self, log: bool) -> Value {
        let ended = self.state.lock().map(|s| s.clone()).unwrap_or_default();
        let mut v = json!({
            "id": self.id,
            "project": self.project,
            "args": self.args,
            "port": self.port,
            "started_ms": self.started_ms,
            "ended_ms": ended.at_ms,
            "exit_code": ended.code,
            "running": ended.at_ms.is_none(),
        });
        if log && let Ok(lines) = self.log.lock() {
            v["log"] = lines.iter().cloned().collect::<Vec<_>>().into();
        }
        v
    }

    fn push_log(&self, line: String) {
        if let Ok(mut l) = self.log.lock() {
            if l.len() == LOG_LINES {
                l.pop_front();
            }
            l.push_back(line);
        }
    }

    /// Forward a request to the run's API: `path` below `/api/`, the run's token added.
    pub async fn proxy(&self, method: Method, path: &str, body: Body) -> Response {
        match self.forward(method, path, body).await {
            Ok(r) => r,
            Err(e) => (
                StatusCode::BAD_GATEWAY,
                axum::Json(json!({ "error": format!("run {}: {e}", self.id) })),
            )
                .into_response(),
        }
    }

    async fn forward(&self, method: Method, path: &str, body: Body) -> Result<Response, String> {
        let stream = tokio::net::TcpStream::connect(("127.0.0.1", self.port))
            .await
            .map_err(|e| format!("engine API not reachable: {e}"))?;
        let (mut sender, conn) = hyper::client::conn::http1::handshake(TokioIo::new(stream))
            .await
            .map_err(|e| e.to_string())?;
        tokio::spawn(async move {
            let _ = conn.await;
        });
        let auth =
            HeaderValue::from_str(&format!("Bearer {}", self.token)).map_err(|e| e.to_string())?;
        let req = Request::builder()
            .method(method)
            .uri(format!("/api/{path}"))
            .header(header::HOST, "127.0.0.1")
            .header(header::AUTHORIZATION, auth)
            .header(header::CONTENT_TYPE, "application/json")
            .body(body)
            .map_err(|e| e.to_string())?;
        let resp = sender.send_request(req).await.map_err(|e| e.to_string())?;
        Ok(resp.map(Body::new).into_response())
    }
}

/// Every run started by this server.
pub struct Runs {
    exe: PathBuf,
    root: PathBuf,
    next: AtomicU32,
    list: Mutex<Vec<Arc<Launched>>>,
}

/// A loopback port free right now.
fn free_port() -> std::io::Result<u16> {
    Ok(std::net::TcpListener::bind("127.0.0.1:0")?
        .local_addr()?
        .port())
}

impl Runs {
    /// Runs of `exe` in `root`.
    pub fn new(exe: PathBuf, root: PathBuf) -> Self {
        Self {
            exe,
            root,
            next: AtomicU32::new(1),
            list: Mutex::new(Vec::new()),
        }
    }

    /// Every run, newest first.
    pub fn all(&self) -> Vec<Arc<Launched>> {
        self.list
            .lock()
            .map(|l| l.iter().rev().cloned().collect())
            .unwrap_or_default()
    }

    /// The run `id`.
    pub fn get(&self, id: u32) -> Option<Arc<Launched>> {
        self.list.lock().ok()?.iter().find(|r| r.id == id).cloned()
    }

    /// Start a run.
    ///
    /// # Errors
    /// An invalid request or a process that cannot start, as text.
    pub fn start(&self, req: &StartRequest) -> Result<Arc<Launched>, String> {
        let mut args = arguments(&self.root, req)?;
        if let Some(dir) = &req.record_dir {
            std::fs::create_dir_all(self.root.join(dir))
                .map_err(|e| format!("record_dir `{dir}`: {e}"))?;
        }
        let port = free_port().map_err(|e| format!("no free port: {e}"))?;
        let token = taktwerk_http::token::generate().map_err(|e| format!("token: {e}"))?;
        args.extend([
            "--http".into(),
            format!("127.0.0.1:{port}"),
            "--http-writes".into(),
        ]);
        let mut child = Command::new(&self.exe)
            .args(&args)
            .current_dir(&self.root)
            .env(taktwerk_http::token::ENV, &token)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .kill_on_drop(true)
            .spawn()
            .map_err(|e| format!("cannot start {}: {e}", self.exe.display()))?;
        let (kill, killed) = oneshot::channel();
        let run = Arc::new(Launched {
            id: self.next.fetch_add(1, Ordering::Relaxed),
            project: req.project.clone(),
            args: args.iter().map(|a| self.shorten(a)).collect(),
            port,
            token,
            started_ms: unix_ms(),
            state: Mutex::new(Ended::default()),
            log: Mutex::new(VecDeque::new()),
            kill: Mutex::new(Some(kill)),
        });
        if let Some(out) = child.stdout.take() {
            tokio::spawn(pump(Arc::clone(&run), BufReader::new(out)));
        }
        if let Some(err) = child.stderr.take() {
            tokio::spawn(pump(Arc::clone(&run), BufReader::new(err)));
        }
        let watched = Arc::clone(&run);
        tokio::spawn(async move {
            let status = tokio::select! {
                s = child.wait() => s,
                _ = killed => {
                    let _ = child.start_kill();
                    child.wait().await
                }
            };
            let code = status.ok().and_then(|s| s.code());
            if let Ok(mut s) = watched.state.lock() {
                *s = Ended {
                    at_ms: Some(unix_ms()),
                    code,
                };
            }
        });
        if let Ok(mut l) = self.list.lock() {
            l.push(Arc::clone(&run));
        }
        Ok(run)
    }

    /// A path argument relative to the folder, for display.
    fn shorten(&self, arg: &str) -> String {
        let p = Path::new(arg);
        if p.is_absolute() && p.starts_with(&self.root) {
            relative(&self.root, p)
        } else {
            arg.to_owned()
        }
    }

    /// Ask the run to stop over its API; kill it when it has not ended after a grace period.
    pub async fn stop(&self, run: &Arc<Launched>) {
        if !run.running() {
            return;
        }
        let resp = run.proxy(Method::POST, "stop", Body::empty()).await;
        if !resp.status().is_success() {
            warn!(run = run.id, status = %resp.status(), "stop over the API failed; killing");
        }
        let deadline = tokio::time::Instant::now() + STOP_GRACE;
        while run.running() && tokio::time::Instant::now() < deadline {
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
        if run.running()
            && let Some(kill) = run.kill.lock().ok().and_then(|mut k| k.take())
        {
            let _ = kill.send(());
        }
    }
}

async fn pump<R: tokio::io::AsyncRead + Unpin>(run: Arc<Launched>, reader: BufReader<R>) {
    let mut lines = reader.lines();
    while let Ok(Some(line)) = lines.next_line().await {
        run.push_log(line);
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, reason = "tests")]
mod tests {
    use super::*;

    #[test]
    fn requests_become_run_arguments() {
        let root = std::env::temp_dir().join(format!("tw-args-{}", std::process::id()));
        std::fs::create_dir_all(&root).unwrap();
        std::fs::write(root.join("p.toml"), "").unwrap();
        std::fs::write(root.join("s.toml"), "").unwrap();
        let root = root.canonicalize().unwrap();
        let p = root.join("p.toml").display().to_string();
        let req = StartRequest {
            project: "p.toml".into(),
            scenario: Some("s.toml".into()),
            speed: Some("max".into()),
            cycles: Some(10),
            set: vec!["a.k=2".into()],
            record_dir: Some("records/x".into()),
            ..StartRequest::default()
        };
        let args = arguments(&root, &req).unwrap();
        assert_eq!(args[..3], ["run".to_owned(), p, "--no-connectors".into()]);
        assert!(args.windows(2).any(|w| w == ["--speed", "max"]));
        assert!(args.windows(2).any(|w| w == ["--set", "a.k=2"]));
        assert!(
            args.windows(2)
                .any(|w| w[0] == "--record-dir" && w[1].ends_with("records/x"))
        );
        let bad = |r: StartRequest| arguments(&root, &r).unwrap_err();
        let base = || StartRequest {
            project: "p.toml".into(),
            ..StartRequest::default()
        };
        assert!(
            bad(StartRequest {
                project: "../p.toml".into(),
                ..base()
            })
            .contains("project")
        );
        assert!(
            bad(StartRequest {
                speed: Some("fast".into()),
                ..base()
            })
            .contains("speed")
        );
        assert!(
            bad(StartRequest {
                set: vec!["k".into()],
                ..base()
            })
            .contains("KEY=VALUE")
        );
        assert!(
            bad(StartRequest {
                record_dir: Some("../x".into()),
                ..base()
            })
            .contains("plain path")
        );
        let with = arguments(
            &root,
            &StartRequest {
                connectors: true,
                ..base()
            },
        )
        .unwrap();
        assert!(!with.contains(&"--no-connectors".to_owned()));
        std::fs::remove_dir_all(&root).unwrap();
    }
}