huskies/server/src/trampoline.rs

//! Detached trampoline — kills the running gateway and launches the replacement.
//!
//! The trampoline is invoked as `huskies --trampoline <job-file>`.  It is spawned
//! as a new Unix session (`setsid`) so that SIGKILL/SIGTERM sent to the original
//! bash-tool process group does not reach it.
//!
//! Flow:
//! 1. Gateway writes a [`TrampolineJob`] atomically and spawns the trampoline.
//! 2. Trampoline backs up the old binary, kills the gateway, starts the new binary.
//! 3. If the new binary passes a health-poll within 10 s → exit 0.
//! 4. If it fails → restore backup, start it with `HUSKIES_TRAMPOLINE_FAILURE` set.

use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::time::Duration;

// ── Job descriptor ────────────────────────────────────────────────────────────

/// Descriptor atomically written by the gateway before it hands control to the trampoline.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrampolineJob {
    /// PID of the currently running gateway process to kill.
    pub gateway_pid: u32,
    /// Absolute path to the newly compiled binary to launch.
    pub new_binary_path: PathBuf,
    /// Absolute path of the binary currently running as the gateway (for rollback).
    pub old_binary_path: PathBuf,
    /// Where to write the backup of the old binary before killing the gateway.
    pub backup_binary_path: PathBuf,
    /// Arguments forwarded verbatim to the new/backup gateway (everything after argv[0]).
    pub gateway_args: Vec<String>,
    /// HTTP URL the trampoline polls to verify the new gateway is serving.
    /// Empty string means skip health polling (used in tests).
    pub health_url: String,
}

// ── Atomic write ──────────────────────────────────────────────────────────────

/// Write `job` to `path` atomically: write to a sibling `.tmp` file, then rename.
///
/// The rename is atomic on POSIX so the trampoline never reads a half-written file.
pub fn write_job_atomic(job: &TrampolineJob, path: &Path) -> Result<(), String> {
    let tmp = path.with_extension("tmp");
    let data = serde_json::to_vec(job).map_err(|e| format!("JSON encode failed: {e}"))?;
    std::fs::write(&tmp, &data).map_err(|e| format!("tmp write failed: {e}"))?;
    std::fs::rename(&tmp, path).map_err(|e| format!("rename failed: {e}"))?;
    Ok(())
}

// ── Spawn detached ────────────────────────────────────────────────────────────

/// Spawn `exe --trampoline <job_path>` as a fully detached process.
///
/// On Unix the child calls `setsid()` in `pre_exec` so it belongs to a new session
/// and is unreachable by signals sent to the original process group.  stdin/stdout/
/// stderr are all redirected to `/dev/null` so the child is fully daemonised.
pub fn spawn_detached_trampoline(exe: &Path, job_path: &Path) -> Result<(), String> {
    let mut cmd = std::process::Command::new(exe);
    cmd.arg("--trampoline").arg(job_path);
    cmd.stdin(std::process::Stdio::null());
    cmd.stdout(std::process::Stdio::null());
    cmd.stderr(std::process::Stdio::null());

    #[cfg(unix)]
    {
        use std::os::unix::process::CommandExt;
        // SAFETY: setsid() is async-signal-safe.  This is called in the child
        // between fork and exec with no other threads running in the child's
        // address space — the only safe window for pre_exec hooks.
        unsafe {
            cmd.pre_exec(|| {
                if libc::setsid() == -1 {
                    return Err(std::io::Error::last_os_error());
                }
                Ok(())
            });
        }
    }

    cmd.spawn().map_err(|e| format!("spawn failed: {e}"))?;
    Ok(())
}

// ── Process management ────────────────────────────────────────────────────────

/// Send SIGTERM to `pid`, wait up to 3 s for it to exit, then SIGKILL.
///
/// After SIGKILL the process is unconditionally considered gone — SIGKILL cannot
/// be ignored, so the process is dead even if it briefly lingers as a zombie
/// (zombie detection via `kill(pid, 0)` is unreliable from a non-parent process).
#[cfg(unix)]
fn kill_gateway_process(pid: u32) -> Result<(), String> {
    use std::thread::sleep;

    let ipid = pid as libc::pid_t;

    // Safety: kill() is always safe to call with any pid.
    let running = || unsafe { libc::kill(ipid, 0) } == 0;

    if !running() {
        return Ok(());
    }

    unsafe { libc::kill(ipid, libc::SIGTERM) };

    for _ in 0..30 {
        sleep(Duration::from_millis(100));
        if !running() {
            return Ok(());
        }
    }

    // SIGKILL cannot be ignored — the kernel will terminate the process.
    // We don't loop-poll after this: the process may briefly appear as a
    // zombie (still in the table, not yet reaped by its parent), in which
    // case kill(pid, 0) returns 0 even though it is effectively dead.
    unsafe { libc::kill(ipid, libc::SIGKILL) };
    sleep(Duration::from_millis(200));
    Ok(())
}

#[cfg(not(unix))]
fn kill_gateway_process(pid: u32) -> Result<(), String> {
    Err(format!("kill not supported on this platform (pid {pid})"))
}

// ── Health polling ────────────────────────────────────────────────────────────

/// Poll `url` every 500 ms until it returns HTTP 2xx or `timeout` elapses.
async fn poll_health(url: &str, timeout: Duration) -> Result<(), String> {
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(2))
        .build()
        .unwrap_or_else(|_| reqwest::Client::new());

    let deadline = std::time::Instant::now() + timeout;
    while std::time::Instant::now() < deadline {
        if let Ok(resp) = client.get(url).send().await
            && resp.status().is_success()
        {
            return Ok(());
        }
        tokio::time::sleep(Duration::from_millis(500)).await;
    }
    Err(format!(
        "health check timed out after {}s: {url}",
        timeout.as_secs()
    ))
}

// ── Core logic (testable) ─────────────────────────────────────────────────────

/// Kill the old gateway, start the new one, and poll its health endpoint.
///
/// Returns `Ok(())` on success or `Err(reason)` when the new gateway could not
/// be started or failed health checks.  Callers are responsible for rollback.
///
/// When `job.health_url` is empty the health poll is skipped (for unit tests).
pub async fn execute_trampoline_core(job: &TrampolineJob) -> Result<(), String> {
    // Back up old binary (best-effort — rollback won't work if this fails).
    if let Some(parent) = job.backup_binary_path.parent() {
        let _ = std::fs::create_dir_all(parent);
    }
    let _ = std::fs::copy(&job.old_binary_path, &job.backup_binary_path);

    // Kill old gateway.  Killing the process closes its file descriptors,
    // which releases the exclusive flock held on `$HOME/.huskies/gateway.pid`.
    // The new gateway (spawned below) will then acquire that flock on startup,
    // ensuring the one-active-gateway invariant is maintained across the swap.
    kill_gateway_process(job.gateway_pid)?;

    // Start new gateway.
    std::process::Command::new(&job.new_binary_path)
        .args(&job.gateway_args)
        .env("HUSKIES_TRAMPOLINE_STARTED", "1")
        .stdin(std::process::Stdio::null())
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .spawn()
        .map_err(|e| format!("failed to start new gateway: {e}"))?;

    // Poll health (skip when URL is empty — used in tests).
    if !job.health_url.is_empty() {
        poll_health(&job.health_url, Duration::from_secs(10)).await?;
    }

    Ok(())
}

// ── Entry point ───────────────────────────────────────────────────────────────

/// Run the trampoline from a job file.  This function never returns.
///
/// On success exits 0 (new gateway is up and will post its own "ready" message).
/// On failure starts the backup binary with `HUSKIES_TRAMPOLINE_FAILURE` set and
/// exits 1.  On unrecoverable failure (cannot start backup either) exits 2.
pub async fn run_trampoline(job_path: &Path) -> ! {
    let data = match std::fs::read(job_path) {
        Ok(d) => d,
        Err(e) => {
            eprintln!(
                "[trampoline] cannot read job file {}: {e}",
                job_path.display()
            );
            std::process::exit(1);
        }
    };

    let job: TrampolineJob = match serde_json::from_slice(&data) {
        Ok(j) => j,
        Err(e) => {
            eprintln!("[trampoline] cannot parse job file: {e}");
            std::process::exit(1);
        }
    };

    eprintln!(
        "[trampoline] killing gateway PID {} and starting {}",
        job.gateway_pid,
        job.new_binary_path.display()
    );

    match execute_trampoline_core(&job).await {
        Ok(()) => {
            eprintln!("[trampoline] new gateway is up — exiting");
            let _ = std::fs::remove_file(job_path);
            std::process::exit(0);
        }
        Err(reason) => {
            eprintln!("[trampoline] new gateway failed ({reason}) — rolling back");
            let result = std::process::Command::new(&job.backup_binary_path)
                .args(&job.gateway_args)
                .env("HUSKIES_TRAMPOLINE_FAILURE", &reason)
                .stdin(std::process::Stdio::null())
                .stdout(std::process::Stdio::null())
                .stderr(std::process::Stdio::null())
                .spawn();
            match result {
                Ok(_) => {
                    let _ = std::fs::remove_file(job_path);
                    std::process::exit(1);
                }
                Err(e) => {
                    eprintln!("[trampoline] FATAL: cannot start backup gateway: {e}");
                    std::process::exit(2);
                }
            }
        }
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    /// Locate `sleep` on the current platform (needed for a portable fake-gateway).
    fn find_sleep() -> PathBuf {
        for candidate in ["/usr/bin/sleep", "/bin/sleep"] {
            let p = PathBuf::from(candidate);
            if p.exists() {
                return p;
            }
        }
        panic!("sleep binary not found");
    }

    #[test]
    fn write_job_atomic_round_trips() {
        let tmp = tempfile::tempdir().unwrap();
        let job = TrampolineJob {
            gateway_pid: 12345,
            new_binary_path: PathBuf::from("/new/huskies"),
            old_binary_path: PathBuf::from("/old/huskies"),
            backup_binary_path: tmp.path().join("backup"),
            gateway_args: vec!["--gateway".to_string(), "/workspace".to_string()],
            health_url: "http://127.0.0.1:3000/api/gateway".to_string(),
        };
        let path = tmp.path().join("trampoline.json");
        write_job_atomic(&job, &path).unwrap();

        // No .tmp file should remain.
        assert!(!path.with_extension("tmp").exists());
        // Final file must exist.
        assert!(path.exists());

        // Round-trip: deserialise and compare fields.
        let data = std::fs::read(&path).unwrap();
        let loaded: TrampolineJob = serde_json::from_slice(&data).unwrap();
        assert_eq!(loaded.gateway_pid, job.gateway_pid);
        assert_eq!(loaded.new_binary_path, job.new_binary_path);
        assert_eq!(loaded.gateway_args, job.gateway_args);
    }

    /// AC 5: a fake-gateway `sleep` process is killed and replaced within timeout.
    #[tokio::test]
    async fn fake_gateway_killed_and_replaced_within_timeout() {
        let sleep_exe = find_sleep();
        let tmp = tempfile::tempdir().unwrap();

        // Spawn the fake gateway (a long-lived sleep process).
        let mut fake_gw = std::process::Command::new(&sleep_exe)
            .arg("60")
            .stdin(std::process::Stdio::null())
            .spawn()
            .expect("spawn fake gateway");
        let fake_pid = fake_gw.id();

        let job = TrampolineJob {
            gateway_pid: fake_pid,
            new_binary_path: sleep_exe.clone(),
            old_binary_path: sleep_exe.clone(),
            backup_binary_path: tmp.path().join("backup"),
            gateway_args: vec!["1".to_string()],
            health_url: String::new(), // skip health check in test
        };

        let start = std::time::Instant::now();
        let result = execute_trampoline_core(&job).await;
        let elapsed = start.elapsed();

        assert!(result.is_ok(), "trampoline core should succeed: {result:?}");
        assert!(
            elapsed < Duration::from_secs(10),
            "should complete well within 10s timeout, took {elapsed:?}"
        );

        // Reap the zombie — should be dead now.
        let status = fake_gw.try_wait().expect("try_wait");
        assert!(
            status.is_some(),
            "fake gateway process should be dead after trampoline kill"
        );
    }
}