xy/crates/xy-supervisor/src/child.rs

use std::sync::Arc;
use tokio::sync::{Mutex, oneshot};

#[async_trait::async_trait]
pub trait ChildHandle: Send + 'static {
    fn pid(&self) -> u32;
    async fn wait(&mut self) -> std::io::Result<Option<i32>>;
    fn terminate(&mut self) -> std::io::Result<()>;
    fn kill(&mut self) -> std::io::Result<()>;
}

pub struct MockChild {
    pid: u32,
    exit_rx: Arc<Mutex<oneshot::Receiver<Option<i32>>>>,
    terminate_tx: Option<oneshot::Sender<()>>,
    kill_tx: Option<oneshot::Sender<()>>,
}

pub struct MockChildController {
    pub exit_tx: Option<oneshot::Sender<Option<i32>>>,
    pub terminate_rx: oneshot::Receiver<()>,
    pub kill_rx: oneshot::Receiver<()>,
}

impl MockChild {
    pub fn new(pid: u32) -> (Self, MockChildController) {
        let (exit_tx, exit_rx) = oneshot::channel();
        let (terminate_tx, terminate_rx) = oneshot::channel();
        let (kill_tx, kill_rx) = oneshot::channel();
        let child = Self {
            pid,
            exit_rx: Arc::new(Mutex::new(exit_rx)),
            terminate_tx: Some(terminate_tx),
            kill_tx: Some(kill_tx),
        };
        let ctl = MockChildController {
            exit_tx: Some(exit_tx),
            terminate_rx,
            kill_rx,
        };
        (child, ctl)
    }
}

#[async_trait::async_trait]
impl ChildHandle for MockChild {
    fn pid(&self) -> u32 {
        self.pid
    }

    async fn wait(&mut self) -> std::io::Result<Option<i32>> {
        let mut rx = self.exit_rx.lock().await;
        match (&mut *rx).await {
            Ok(code) => Ok(code),
            Err(_) => Err(std::io::Error::other("exit_tx dropped")),
        }
    }

    fn terminate(&mut self) -> std::io::Result<()> {
        if let Some(tx) = self.terminate_tx.take() {
            let _ = tx.send(());
        }
        Ok(())
    }

    fn kill(&mut self) -> std::io::Result<()> {
        if let Some(tx) = self.kill_tx.take() {
            let _ = tx.send(());
        }
        Ok(())
    }
}

use crate::logs::LogSink;
use nix::sys::signal::{Signal, kill};
use nix::unistd::Pid;
use std::process::Stdio;
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::{Child as TokioChild, Command};
use xy_protocol::{ServerConfig, rpc::LogStream};

pub struct RealChild {
    pid: u32,
    pgid: Pid,
    child: Option<TokioChild>,
}

impl RealChild {
    pub fn pgid(&self) -> Pid {
        self.pgid
    }
}

#[async_trait::async_trait]
impl ChildHandle for RealChild {
    fn pid(&self) -> u32 {
        self.pid
    }

    async fn wait(&mut self) -> std::io::Result<Option<i32>> {
        let child = self
            .child
            .as_mut()
            .ok_or_else(|| std::io::Error::other("already waited"))?;

        let status = child.wait().await?;

        Ok(status.code())
    }

    fn terminate(&mut self) -> std::io::Result<()> {
        kill(Pid::from_raw(-self.pgid.as_raw()), Signal::SIGTERM)
            .map_err(|err| std::io::Error::other(err.to_string()))
    }

    fn kill(&mut self) -> std::io::Result<()> {
        kill(Pid::from_raw(-self.pgid.as_raw()), Signal::SIGKILL)
            .map_err(|err| std::io::Error::other(err.to_string()))
    }
}

pub fn spawn_with_logs(cfg: &ServerConfig, sink: LogSink) -> std::io::Result<RealChild> {
    let mut cmd = Command::new(&cfg.command);

    cmd.args(&cfg.args);

    for (k, v) in &cfg.env {
        cmd.env(k, v);
    }

    if let Some(dir) = &cfg.working_dir {
        cmd.current_dir(dir);
    }

    cmd.stdout(Stdio::piped());
    cmd.stderr(Stdio::piped());
    cmd.kill_on_drop(true);

    // Own process group so signals reach the whole tree.
    unsafe {
        cmd.pre_exec(|| {
            nix::unistd::setpgid(Pid::from_raw(0), Pid::from_raw(0))
                .map_err(|err| std::io::Error::other(err.to_string()))
        });
    }

    let mut child = cmd.spawn()?;

    let pid = child.id().ok_or_else(|| std::io::Error::other("no pid"))?;
    let pgid = Pid::from_raw(pid as i32);

    if let Some(out) = child.stdout.take() {
        spawn_pump(out, sink.clone(), LogStream::Stdout);
    }

    if let Some(err) = child.stderr.take() {
        spawn_pump(err, sink.clone(), LogStream::Stderr);
    }

    Ok(RealChild {
        pid,
        pgid,
        child: Some(child),
    })
}

fn spawn_pump<R: tokio::io::AsyncRead + Unpin + Send + 'static>(
    reader: R,
    sink: LogSink,
    stream: LogStream,
) {
    tokio::spawn(async move {
        let mut lines = BufReader::new(reader).lines();

        loop {
            match lines.next_line().await {
                Ok(Some(line)) => sink.record(stream, line),
                Ok(None) => break,
                Err(err) => {
                    tracing::warn!(
                        server = %sink.server_name,
                        error = %err,
                        ?stream,
                        "log pump read error"
                    );
                    break;
                }
            }
        }
    });
}

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

    #[tokio::test]
    async fn mock_child_exit() {
        let (mut child, mut ctl) = MockChild::new(123);
        assert_eq!(child.pid(), 123);
        ctl.exit_tx.take().unwrap().send(Some(0)).unwrap();
        assert_eq!(child.wait().await.unwrap(), Some(0));
    }

    #[tokio::test]
    async fn mock_child_terminate() {
        let (mut child, mut ctl) = MockChild::new(1);
        child.terminate().unwrap();
        ctl.terminate_rx.try_recv().unwrap();
    }
}