clusterflux-public/crates/disasmer-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
            // descendants such as rootless Podman and the command inside its container.
            unsafe {
                libc::kill(process_group, libc::SIGKILL);
            }
        }
        let _ = child.kill();
        let _ = child.wait();
    }

    #[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> {
        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_process_group(&mut child);
                return Err(BackendError::Command(format!(
                    "establish execution control channel: {error}"
                )));
            }
        };

        let mut native_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_process_group(&mut child);
                    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_process_group(&mut child);
                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_process_group(&mut child);
                    let _ = stdout.join();
                    let _ = stderr.join();
                    return Err(BackendError::Cancelled(
                        "coordinator requested cancellation or abort".to_owned(),
                    ));
                }
                Ok(false) => {}
                Err(error) => {
                    Self::terminate_process_group(&mut child);
                    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 native_frozen_epoch == Some(epoch) {
                        #[cfg(unix)]
                        Self::resume_process_group(&child)?;
                        self.set_command_status(format!(
                            "running native command pid {} after debug epoch {epoch} resumed",
                            child.id()
                        ));
                        self.debug_control.mark_running(epoch);
                        native_frozen_epoch = None;
                    }
                } else if native_frozen_epoch != Some(epoch)
                    && self.debug_control.frozen_epoch() != Some(epoch)
                {
                    #[cfg(unix)]
                    {
                        Self::freeze_process_group(&child)?;
                        self.set_command_status(format!(
                            "frozen native command pid {} for debug epoch {epoch}",
                            child.id()
                        ));
                        self.debug_control.mark_frozen(epoch);
                        native_frozen_epoch = Some(epoch);
                    }
                }
            }
            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,
        })
    }
}