clusterflux-public/crates/clusterflux-node/src/command_runner.rs

use clusterflux_core::{
    CommandInvocation, Digest, LogBuffer, NativeCommandPolicy, NodeId, TaskInstanceId, VfsObject,
    VfsOverlay, VfsPath,
};
use serde::{Deserialize, Serialize};

use super::BackendError;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CapturedCommandLogs {
    pub stdout: String,
    pub stderr: String,
    pub stdout_truncated: bool,
    pub stderr_truncated: bool,
    pub backpressured: bool,
}

pub const DEFAULT_COMMAND_LOG_LIMIT_BYTES: usize = 256 * 1024;

pub fn authorize_node_command(
    hosted_control_plane: bool,
    node_has_command_capability: bool,
) -> Result<(), BackendError> {
    NativeCommandPolicy {
        hosted_control_plane,
        node_has_command_capability,
    }
    .authorize()
    .map_err(BackendError::Denied)
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct VirtualThreadCommand {
    pub virtual_thread: TaskInstanceId,
    pub invocation: CommandInvocation,
    pub stage_stdout_as: Option<VfsPath>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CommandOutput {
    pub virtual_thread: TaskInstanceId,
    pub status_code: Option<i32>,
    pub stdout: String,
    pub stderr: String,
    pub stdout_truncated: bool,
    pub stderr_truncated: bool,
    pub log_backpressured: bool,
    pub staged_artifact: Option<VfsObject>,
}

#[derive(Clone, Debug)]
pub struct LocalCommandExecutor {
    pub node: NodeId,
    pub hosted_control_plane: bool,
    pub has_command_capability: bool,
}

impl LocalCommandExecutor {
    pub fn run(
        &self,
        command: VirtualThreadCommand,
        overlay: &mut VfsOverlay,
    ) -> Result<CommandOutput, BackendError> {
        self.run_with_log_limit(command, overlay, DEFAULT_COMMAND_LOG_LIMIT_BYTES)
    }

    pub fn run_with_log_limit(
        &self,
        command: VirtualThreadCommand,
        overlay: &mut VfsOverlay,
        max_log_bytes: usize,
    ) -> Result<CommandOutput, BackendError> {
        authorize_node_command(self.hosted_control_plane, self.has_command_capability)?;

        let output = std::process::Command::new(&command.invocation.program)
            .args(&command.invocation.args)
            .output()
            .map_err(|err| BackendError::Command(format!("{err:#}")))?;

        let logs = capture_command_logs(
            &command.virtual_thread,
            &output.stdout,
            &output.stderr,
            max_log_bytes,
        );
        let staged_artifact = if let Some(path) = command.stage_stdout_as {
            Some(overlay.write(
                path,
                Digest::sha256(&output.stdout),
                output.stdout.len() as u64,
            ))
        } else {
            None
        };

        Ok(CommandOutput {
            virtual_thread: command.virtual_thread,
            status_code: output.status.code(),
            stdout: logs.stdout,
            stderr: logs.stderr,
            stdout_truncated: logs.stdout_truncated,
            stderr_truncated: logs.stderr_truncated,
            log_backpressured: logs.backpressured,
            staged_artifact,
        })
    }
}

pub(super) fn capture_command_logs(
    task: &TaskInstanceId,
    stdout: &[u8],
    stderr: &[u8],
    max_log_bytes: usize,
) -> CapturedCommandLogs {
    let mut logs = LogBuffer::new(max_log_bytes);
    logs.push(task.clone(), stdout);
    logs.push(task.clone(), stderr);
    let records = logs.records();
    let stdout_record = &records[0];
    let stderr_record = &records[1];
    debug_assert_eq!(&stdout_record.task, task);
    debug_assert_eq!(&stderr_record.task, task);
    CapturedCommandLogs {
        stdout: String::from_utf8_lossy(&stdout_record.bytes).into_owned(),
        stderr: String::from_utf8_lossy(&stderr_record.bytes).into_owned(),
        stdout_truncated: stdout_record.truncated,
        stderr_truncated: stderr_record.truncated,
        backpressured: logs.backpressured(),
    }
}