clusterflux-public/crates/clusterflux-node/src/assignment_runner/process_runner.rs

use super::*;

pub(super) struct CoordinatorControlledProcessRunner {
    pub(super) args: Args,
    pub(super) process: String,
    pub(super) task: String,
    pub(super) node_private_key: String,
    pub(super) debug_control: Arc<WasmDebugControl>,
    pub(super) command_status: Arc<Mutex<Option<String>>>,
    pub(super) timeout: Duration,
}

impl CoordinatorControlledProcessRunner {
    const MAX_CAPTURE_BYTES: usize = 256 * 1024 + 1;

    pub(super) fn new(host: &CoordinatorWasmTaskHost, timeout: Duration) -> Self {
        Self {
            args: host.args.clone(),
            process: host.process.clone(),
            task: host.parent_task.clone(),
            node_private_key: host.node_private_key.clone(),
            debug_control: Arc::clone(&host.debug_control),
            command_status: Arc::clone(&host.command_status),
            timeout,
        }
    }

    fn set_command_status(&self, status: impl Into<String>) {
        if let Ok(mut current) = self.command_status.lock() {
            *current = Some(status.into());
        }
    }

    fn abort_requested(&self, session: &mut CoordinatorSession) -> Result<bool, BackendError> {
        let request = signed_node_request_json(
            &self.args,
            &self.node_private_key,
            "poll_task_control",
            serde_json::json!({
                "type": "poll_task_control",
                "tenant": &self.args.tenant,
                "project": &self.args.project,
                "process": &self.process,
                "node": &self.args.node,
                "task": &self.task,
            }),
        )
        .map_err(|error| BackendError::Command(error.to_string()))?;
        let response = session
            .request(request)
            .map_err(|error| BackendError::Command(format!("poll task control: {error}")))?;
        Ok(response
            .get("abort_requested")
            .and_then(serde_json::Value::as_bool)
            .unwrap_or(false))
    }

    fn terminate_process_group(child: &mut std::process::Child) {
        #[cfg(unix)]
        {
            let process_group = -(child.id() as i32);
            // The child is placed in a new process group before exec, so this reaches
            // all native descendants. Podman containers are removed separately.
            unsafe {
                libc::kill(process_group, libc::SIGKILL);
            }
        }
        let _ = child.kill();
        let _ = child.wait();
    }

    fn podman_container_name(command: &PodmanCommand) -> Option<String> {
        if command.program != "podman" || command.args.first().map(String::as_str) != Some("run") {
            return None;
        }
        command
            .args
            .windows(2)
            .find(|arguments| arguments[0] == "--name")
            .map(|arguments| arguments[1].clone())
    }

    fn set_podman_paused(container: &str, paused: bool) -> Result<(), BackendError> {
        let action = if paused { "pause" } else { "unpause" };
        let output = std::process::Command::new("podman")
            .args([action, container])
            .output()
            .map_err(|error| {
                BackendError::Command(format!(
                    "failed to invoke `podman {action}` for container `{container}`: {error}"
                ))
            })?;
        if !output.status.success() {
            return Err(BackendError::Command(format!(
                "`podman {action}` failed for container `{container}`: {}",
                String::from_utf8_lossy(&output.stderr).trim()
            )));
        }

        let inspection = std::process::Command::new("podman")
            .args(["inspect", "--format", "{{.State.Paused}}", container])
            .output()
            .map_err(|error| {
                BackendError::Command(format!(
                    "failed to inspect Podman pause state for container `{container}`: {error}"
                ))
            })?;
        let observed = String::from_utf8_lossy(&inspection.stdout);
        let expected = if paused { "true" } else { "false" };
        if !inspection.status.success() || observed.trim() != expected {
            return Err(BackendError::Command(format!(
                "container `{container}` did not verify as {} after `podman {action}`: status={:?} stdout={} stderr={}",
                if paused { "paused" } else { "running" },
                inspection.status.code(),
                observed.trim(),
                String::from_utf8_lossy(&inspection.stderr).trim()
            )));
        }
        Ok(())
    }

    fn terminate_execution(child: &mut std::process::Child, container: Option<&str>) {
        if let Some(container) = container {
            let _ = std::process::Command::new("podman")
                .args(["rm", "--force", container])
                .output();
        }
        Self::terminate_process_group(child);
    }

    fn freeze_execution(
        child: &std::process::Child,
        container: Option<&str>,
    ) -> Result<(), BackendError> {
        if let Some(container) = container {
            return Self::set_podman_paused(container, true);
        }
        #[cfg(unix)]
        {
            Self::freeze_process_group(child)
        }
        #[cfg(not(unix))]
        {
            let _ = child;
            Err(BackendError::Command(
                "native debug freeze requires Unix process groups".to_owned(),
            ))
        }
    }

    fn resume_execution(
        child: &std::process::Child,
        container: Option<&str>,
    ) -> Result<(), BackendError> {
        if let Some(container) = container {
            return Self::set_podman_paused(container, false);
        }
        #[cfg(unix)]
        {
            Self::resume_process_group(child)
        }
        #[cfg(not(unix))]
        {
            let _ = child;
            Err(BackendError::Command(
                "native debug resume requires Unix process groups".to_owned(),
            ))
        }
    }

    #[cfg(unix)]
    fn freeze_process_group(child: &std::process::Child) -> Result<(), BackendError> {
        let process_group = -(child.id() as i32);
        let result = unsafe { libc::kill(process_group, libc::SIGSTOP) };
        if result == 0 {
            Ok(())
        } else {
            Err(BackendError::Command(format!(
                "failed to freeze native process group {}: {}",
                child.id(),
                std::io::Error::last_os_error()
            )))
        }
    }

    #[cfg(unix)]
    fn resume_process_group(child: &std::process::Child) -> Result<(), BackendError> {
        let process_group = -(child.id() as i32);
        let result = unsafe { libc::kill(process_group, libc::SIGCONT) };
        if result == 0 {
            Ok(())
        } else {
            Err(BackendError::Command(format!(
                "failed to resume native process group {}: {}",
                child.id(),
                std::io::Error::last_os_error()
            )))
        }
    }

    fn drain_bounded(
        mut reader: impl Read + Send + 'static,
        maximum: usize,
    ) -> thread::JoinHandle<Result<Vec<u8>, String>> {
        thread::spawn(move || {
            let mut captured = Vec::new();
            let mut buffer = [0_u8; 16 * 1024];
            loop {
                let count = reader
                    .read(&mut buffer)
                    .map_err(|error| error.to_string())?;
                if count == 0 {
                    break;
                }
                let remaining = maximum.saturating_sub(captured.len());
                captured.extend_from_slice(&buffer[..count.min(remaining)]);
            }
            Ok(captured)
        })
    }
}

impl ProcessRunner for CoordinatorControlledProcessRunner {
    fn run(&mut self, command: &PodmanCommand) -> Result<ProcessOutput, BackendError> {
        let podman_container = Self::podman_container_name(command);
        self.set_command_status(format!(
            "starting native command: {} {}",
            command.program,
            command.args.join(" ")
        ));
        let mut process = std::process::Command::new(&command.program);
        process
            .args(&command.args)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        #[cfg(unix)]
        process.process_group(0);
        let mut child = process
            .spawn()
            .map_err(|error| BackendError::Command(error.to_string()))?;
        self.set_command_status(format!(
            "running native command pid {}: {} {}",
            child.id(),
            command.program,
            command.args.join(" ")
        ));
        let stdout = Self::drain_bounded(
            child
                .stdout
                .take()
                .ok_or_else(|| BackendError::Command("command stdout pipe missing".to_owned()))?,
            Self::MAX_CAPTURE_BYTES,
        );
        let stderr = Self::drain_bounded(
            child
                .stderr
                .take()
                .ok_or_else(|| BackendError::Command("command stderr pipe missing".to_owned()))?,
            Self::MAX_CAPTURE_BYTES,
        );
        let mut session = match CoordinatorSession::connect(&self.args.coordinator) {
            Ok(session) => session,
            Err(error) => {
                Self::terminate_execution(&mut child, podman_container.as_deref());
                return Err(BackendError::Command(format!(
                    "establish execution control channel: {error}"
                )));
            }
        };

        let mut frozen_epoch = None;
        let started = Instant::now();
        let status = loop {
            match child.try_wait() {
                Ok(Some(status)) => break status,
                Ok(None) => {}
                Err(error) => {
                    Self::terminate_execution(&mut child, podman_container.as_deref());
                    return Err(BackendError::Command(error.to_string()));
                }
            }
            if started.elapsed() >= self.timeout {
                self.set_command_status(format!(
                    "native command pid {} exceeded wall-clock timeout of {} ms",
                    child.id(),
                    self.timeout.as_millis()
                ));
                Self::terminate_execution(&mut child, podman_container.as_deref());
                let _ = stdout.join();
                let _ = stderr.join();
                return Err(BackendError::Command(format!(
                    "native command exceeded wall-clock timeout of {} ms",
                    self.timeout.as_millis()
                )));
            }
            match self.abort_requested(&mut session) {
                Ok(true) => {
                    self.set_command_status(format!(
                        "aborting native command pid {} at coordinator request",
                        child.id()
                    ));
                    Self::terminate_execution(&mut child, podman_container.as_deref());
                    let _ = stdout.join();
                    let _ = stderr.join();
                    return Err(BackendError::Cancelled(
                        "coordinator requested cancellation or abort".to_owned(),
                    ));
                }
                Ok(false) => {}
                Err(error) => {
                    Self::terminate_execution(&mut child, podman_container.as_deref());
                    let _ = stdout.join();
                    let _ = stderr.join();
                    return Err(error);
                }
            }
            if let Some(epoch) = self.debug_control.requested_epoch() {
                if self.debug_control.resume_requested(epoch) {
                    if frozen_epoch == Some(epoch) {
                        match Self::resume_execution(&child, podman_container.as_deref()) {
                            Ok(()) => {
                                self.set_command_status(format!(
                                    "running command pid {} after debug epoch {epoch} resumed",
                                    child.id()
                                ));
                                self.debug_control.mark_running(epoch);
                                frozen_epoch = None;
                            }
                            Err(error) => self.set_command_status(format!(
                                "debug epoch {epoch} resume is pending: {error}"
                            )),
                        }
                    }
                } else if frozen_epoch != Some(epoch)
                    && self.debug_control.frozen_epoch() != Some(epoch)
                {
                    match Self::freeze_execution(&child, podman_container.as_deref()) {
                        Ok(()) => {
                            self.set_command_status(format!(
                                "frozen command pid {} for debug epoch {epoch}",
                                child.id()
                            ));
                            self.debug_control.mark_frozen(epoch);
                            frozen_epoch = Some(epoch);
                        }
                        Err(error) => self.set_command_status(format!(
                            "debug epoch {epoch} freeze is pending: {error}"
                        )),
                    }
                }
            }
            thread::sleep(Duration::from_millis(50));
        };
        let stdout = stdout
            .join()
            .map_err(|_| BackendError::Command("stdout reader panicked".to_owned()))?
            .map_err(BackendError::Command)?;
        let stderr = stderr
            .join()
            .map_err(|_| BackendError::Command("stderr reader panicked".to_owned()))?
            .map_err(BackendError::Command)?;
        self.set_command_status(format!(
            "native command exited with status {:?}",
            status.code()
        ));
        Ok(ProcessOutput {
            status_code: status.code(),
            stdout,
            stderr,
        })
    }
}