maestro/crates/maestro-e2e/src/daemon.rs

//! Subprocess wrapper around `maestrod`.

use std::net::{Ipv4Addr, SocketAddr};
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::time::Duration;

use tempfile::TempDir;

use crate::build::workspace_target_bin;
use crate::util::{free_port, wait_for_async};

/// A live `maestrod` running in a child process. Kill the child on drop.
///
/// `state_dir`, `home`, and `tmux_tmpdir` are kept as `TempDir`s so they
/// outlive the process and get cleaned up automatically.
#[derive(Debug)]
pub struct Daemon {
    pub daemon_id: String,
    pub bind_addr: SocketAddr,
    pub daemon_url: String,
    child: Option<Child>,
    home: TempDir,
    state: TempDir,
    tmux: TempDir,
    registry_url: String,
    token: String,
}

impl Daemon {
    /// Spawn a fresh daemon pointing at `registry_url` (the daemon-facing WS
    /// URL — usually `ws://<addr>/daemon`).
    pub async fn spawn(daemon_id: &str, registry_url: &str, token: &str) -> Self {
        let port = free_port();
        let bind = SocketAddr::from((Ipv4Addr::LOCALHOST, port));
        Self::spawn_with_port(daemon_id, registry_url, token, bind).await
    }

    /// Stop the underlying process, but keep its state dirs around so we can
    /// `restart()` against the same `$HOME` / `state` / `TMUX_TMPDIR`.
    pub fn stop(&mut self) {
        if let Some(mut c) = self.child.take() {
            let _ = c.kill();
            let _ = c.wait();
        }
    }

    /// Re-spawn the daemon on the same port + state. The previous tmux
    /// server (under `TMUX_TMPDIR`) is left running so `recover()` can adopt
    /// its sessions.
    pub async fn restart(&mut self) {
        self.stop();
        let child = launch(
            &self.daemon_id,
            self.bind_addr,
            &self.daemon_url,
            &self.registry_url,
            &self.token,
            self.home.path(),
            self.state.path(),
            self.tmux.path(),
        );
        self.child = Some(child);
        self.wait_until_attach_ready().await;
    }

    async fn spawn_with_port(
        daemon_id: &str,
        registry_url: &str,
        token: &str,
        bind: SocketAddr,
    ) -> Self {
        let home = TempDir::new().expect("home tmpdir");
        let state = TempDir::new().expect("state tmpdir");
        let tmux = TempDir::new().expect("tmux tmpdir");

        let daemon_url = format!("ws://{bind}");
        let child = launch(
            daemon_id,
            bind,
            &daemon_url,
            registry_url,
            token,
            home.path(),
            state.path(),
            tmux.path(),
        );
        let me = Self {
            daemon_id: daemon_id.to_string(),
            bind_addr: bind,
            daemon_url,
            child: Some(child),
            home,
            state,
            tmux,
            registry_url: registry_url.to_string(),
            token: token.to_string(),
        };
        me.wait_until_attach_ready().await;
        me
    }

    async fn wait_until_attach_ready(&self) {
        let addr = self.bind_addr;
        wait_for_async(
            "daemon attach port",
            Duration::from_secs(5),
            || async move { tokio::net::TcpStream::connect(addr).await.is_ok() },
        )
        .await;
    }

    pub fn home(&self) -> &std::path::Path {
        self.home.path()
    }

    pub fn state_dir(&self) -> PathBuf {
        // Matches maestro_daemon::default_state_dir which uses XDG_STATE_HOME.
        self.home.path().join(".local/state/maestrod")
    }
}

impl Drop for Daemon {
    fn drop(&mut self) {
        if let Some(mut c) = self.child.take() {
            let _ = c.kill();
            let _ = c.wait();
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn launch(
    daemon_id: &str,
    bind: SocketAddr,
    daemon_url: &str,
    registry_url: &str,
    token: &str,
    home: &std::path::Path,
    _state: &std::path::Path,
    tmux_tmpdir: &std::path::Path,
) -> Child {
    let bin = workspace_target_bin("maestro-daemon", "maestrod");
    Command::new(&bin)
        .env_clear()
        .env(
            "PATH",
            std::env::var("PATH").unwrap_or_else(|_| "/usr/bin:/bin".into()),
        )
        .env("HOME", home)
        .env("XDG_CONFIG_HOME", home.join(".config"))
        .env("XDG_STATE_HOME", home.join(".local/state"))
        .env("TMUX_TMPDIR", tmux_tmpdir)
        .env("MAESTRO_TOKEN", token)
        .env("MAESTROD_DAEMON_ID", daemon_id)
        .env("MAESTROD_REGISTRY_URL", registry_url)
        .env("MAESTROD_BIND_ADDR", bind.to_string())
        .env("MAESTROD_DAEMON_URL", daemon_url)
        .env("RUST_LOG", "warn")
        .stdout(Stdio::null())
        .stderr(Stdio::inherit())
        .arg("run")
        .spawn()
        .unwrap_or_else(|e| panic!("spawn {} failed: {e}", bin.display()))
}