clusterflux-public/crates/disasmer-node/src/main.rs

use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::path::PathBuf;
use std::time::{Duration, Instant};

use disasmer_core::{
    Capability, CommandInvocation, NodeCapabilities, NodeId, TaskId, VfsOverlay, VfsPath,
};
use disasmer_node::{CommandOutput, LocalCommandExecutor, VirtualThreadCommand};
use serde_json::{json, Value};

#[derive(Debug)]
struct Args {
    coordinator: String,
    tenant: String,
    project: String,
    node: String,
    process: String,
    task: String,
    command: String,
    command_args: Vec<String>,
    artifact: String,
    enrollment_grant: Option<String>,
    public_key: Option<String>,
    control_poll_ms: u64,
    assignment_poll_ms: u64,
    emit_ready: bool,
    worker: bool,
}

#[derive(Clone, Debug)]
struct RuntimeTask {
    process: String,
    task: String,
    command: String,
    command_args: Vec<String>,
    artifact: String,
    epoch: Option<u64>,
    task_assignment_response: Value,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = parse_args()?;

    let mut session = CoordinatorSession::connect(&args.coordinator)?;

    let registration = register_node(&mut session, &args)?;
    let heartbeat = session.request(json!({
        "type": "node_heartbeat",
        "node": &args.node,
    }))?;
    let capability_report = session.request(json!({
        "type": "report_node_capabilities",
        "tenant": &args.tenant,
        "project": &args.project,
        "node": &args.node,
        "capabilities": NodeCapabilities::detect_current(),
        "cached_environment_digests": [],
        "dependency_cache_digests": [],
        "source_snapshots": [],
        "artifact_locations": [],
        "direct_connectivity": true,
        "online": true,
    }))?;

    if args.worker {
        return worker_loop(
            &args,
            &mut session,
            registration,
            heartbeat,
            capability_report,
        );
    }

    let task_assignment = session.request(json!({
        "type": "schedule_task",
        "tenant": &args.tenant,
        "project": &args.project,
        "environment": null,
        "environment_digest": null,
        "required_capabilities": [Capability::Command],
        "dependency_cache": null,
        "source_snapshot": null,
        "required_artifacts": [],
        "quota_available": true,
        "policy_allowed": true,
        "prefer_node": &args.node,
    }))?;
    let runtime_task = RuntimeTask {
        process: args.process.clone(),
        task: args.task.clone(),
        command: args.command.clone(),
        command_args: args.command_args.clone(),
        artifact: args.artifact.clone(),
        epoch: None,
        task_assignment_response: task_assignment,
    };
    let report = run_runtime_task(
        &args,
        &mut session,
        runtime_task,
        registration,
        heartbeat,
        capability_report,
    )?;
    println!("{}", serde_json::to_string(&report)?);
    Ok(())
}

fn worker_loop(
    args: &Args,
    session: &mut CoordinatorSession,
    registration: Value,
    heartbeat: Value,
    capability_report: Value,
) -> Result<(), Box<dyn std::error::Error>> {
    if args.emit_ready {
        println!(
            "{}",
            serde_json::to_string(&json!({
                "node_status": "ready",
                "mode": "worker",
                "node": &args.node,
            }))?
        );
        std::io::stdout().flush()?;
    }

    loop {
        let response = session.request(json!({
            "type": "poll_task_assignment",
            "tenant": &args.tenant,
            "project": &args.project,
            "node": &args.node,
        }))?;
        let Some(assignment) = response.get("assignment").filter(|value| !value.is_null()) else {
            std::thread::sleep(Duration::from_millis(args.assignment_poll_ms));
            continue;
        };
        let runtime_task = runtime_task_from_assignment(assignment)?;
        let report = run_runtime_task(
            args,
            session,
            runtime_task,
            registration.clone(),
            heartbeat.clone(),
            capability_report.clone(),
        )?;
        println!("{}", serde_json::to_string(&report)?);
        std::io::stdout().flush()?;
    }
}

fn runtime_task_from_assignment(value: &Value) -> Result<RuntimeTask, Box<dyn std::error::Error>> {
    Ok(RuntimeTask {
        process: required_string(value, "process")?,
        task: required_string(value, "task")?,
        command: required_string(value, "command")?,
        command_args: value
            .get("command_args")
            .and_then(Value::as_array)
            .map(|items| {
                items
                    .iter()
                    .filter_map(Value::as_str)
                    .map(str::to_owned)
                    .collect::<Vec<_>>()
            })
            .unwrap_or_default(),
        artifact: required_string(value, "artifact_path")?,
        epoch: value.get("epoch").and_then(Value::as_u64),
        task_assignment_response: value.clone(),
    })
}

fn required_string(value: &Value, field: &str) -> Result<String, Box<dyn std::error::Error>> {
    value
        .get(field)
        .and_then(Value::as_str)
        .map(str::to_owned)
        .ok_or_else(|| format!("task assignment missing string field `{field}`").into())
}

fn run_runtime_task(
    args: &Args,
    session: &mut CoordinatorSession,
    task: RuntimeTask,
    registration: Value,
    heartbeat: Value,
    capability_report: Value,
) -> Result<Value, Box<dyn std::error::Error>> {
    let epoch = match task.epoch {
        Some(epoch) => epoch,
        None => {
            let started = session.request(json!({
                "type": "start_process",
                "tenant": &args.tenant,
                "project": &args.project,
                "process": &task.process,
            }))?;
            started
                .get("epoch")
                .and_then(Value::as_u64)
                .ok_or("coordinator start_process response missing epoch")?
        }
    };
    session.request(json!({
        "type": "reconnect_node",
        "node": &args.node,
        "process": &task.process,
        "epoch": epoch,
    }))?;
    let debug_command = session.request(json!({
        "type": "poll_debug_command",
        "tenant": &args.tenant,
        "project": &args.project,
        "process": &task.process,
        "node": &args.node,
        "task": &task.task,
    }))?;
    if args.emit_ready && !args.worker {
        println!(
            "{}",
            serde_json::to_string(&json!({
                "node_status": "ready",
                "node": &args.node,
                "process": &task.process,
                "task": &task.task,
            }))?
        );
        std::io::stdout().flush()?;
    }
    if args.control_poll_ms > 0 && wait_for_cancellation(session, args, &task)? {
        let recorded = session.request(json!({
            "type": "task_completed",
            "tenant": &args.tenant,
            "project": &args.project,
            "process": &task.process,
            "node": &args.node,
            "task": &task.task,
            "terminal_state": "cancelled",
            "status_code": null,
            "stdout_bytes": 0,
            "stderr_bytes": 0,
            "artifact_path": null,
            "artifact_digest": null,
            "artifact_size_bytes": null,
        }))?;
        return Ok(cancelled_node_report(
            args,
            &task,
            registration,
            heartbeat,
            capability_report,
            task.task_assignment_response.clone(),
            debug_command,
            recorded,
            session.requests(),
        ));
    }

    let executor = LocalCommandExecutor {
        node: NodeId::new(args.node.clone()),
        hosted_control_plane: false,
        has_command_capability: true,
    };
    let task_id = TaskId::new(task.task.clone());
    let mut overlay = VfsOverlay::new(task_id.clone(), NodeId::new(args.node.clone()));
    let output = executor.run(
        VirtualThreadCommand {
            virtual_thread: task_id,
            invocation: CommandInvocation {
                program: task.command.clone(),
                args: task.command_args.clone(),
                env: None,
            },
            stage_stdout_as: Some(VfsPath::new(task.artifact.clone())?),
        },
        &mut overlay,
    )?;
    let manifest = overlay.flush();
    let staged = output.staged_artifact.as_ref();
    let artifact_digest = staged.map(|artifact| artifact.digest.clone());
    let artifact_path = staged.map(|artifact| artifact.path.as_str().to_owned());
    let artifact_size_bytes = staged.map(|artifact| artifact.size);

    let log_event = session.request(json!({
        "type": "report_task_log",
        "tenant": &args.tenant,
        "project": &args.project,
        "process": &task.process,
        "node": &args.node,
        "task": &task.task,
        "stdout_bytes": output.stdout.len(),
        "stderr_bytes": output.stderr.len(),
        "stdout_tail": &output.stdout,
        "stderr_tail": &output.stderr,
        "stdout_truncated": output.stdout_truncated,
        "stderr_truncated": output.stderr_truncated,
        "backpressured": output.log_backpressured,
    }))?;
    let vfs_metadata = session.request(json!({
        "type": "report_vfs_metadata",
        "tenant": &args.tenant,
        "project": &args.project,
        "process": &task.process,
        "node": &args.node,
        "task": &task.task,
        "artifact_path": artifact_path,
        "artifact_digest": artifact_digest,
        "artifact_size_bytes": artifact_size_bytes,
        "large_bytes_uploaded": manifest.large_bytes_uploaded,
    }))?;
    let recorded = session.request(json!({
        "type": "task_completed",
        "tenant": &args.tenant,
        "project": &args.project,
        "process": &task.process,
        "node": &args.node,
        "task": &task.task,
        "status_code": output.status_code,
        "stdout_bytes": output.stdout.len(),
        "stderr_bytes": output.stderr.len(),
        "stdout_tail": &output.stdout,
        "stderr_tail": &output.stderr,
        "stdout_truncated": output.stdout_truncated,
        "stderr_truncated": output.stderr_truncated,
        "artifact_path": staged.map(|artifact| artifact.path.as_str().to_owned()),
        "artifact_digest": staged.map(|artifact| artifact.digest.clone()),
        "artifact_size_bytes": artifact_size_bytes,
    }))?;

    Ok(node_report(
        output,
        manifest.large_bytes_uploaded,
        registration,
        heartbeat,
        capability_report,
        task.task_assignment_response,
        debug_command,
        log_event,
        vfs_metadata,
        recorded,
        session.requests(),
    ))
}

fn node_report(
    output: CommandOutput,
    large_bytes_uploaded: bool,
    registration_response: Value,
    heartbeat_response: Value,
    capability_response: Value,
    task_assignment_response: Value,
    debug_command_response: Value,
    log_event_response: Value,
    vfs_metadata_response: Value,
    coordinator_response: Value,
    session_requests: usize,
) -> Value {
    json!({
        "node_status": "completed",
        "virtual_thread": output.virtual_thread,
        "terminal_state": if output.status_code == Some(0) { "completed" } else { "failed" },
        "status_code": output.status_code,
        "stdout_bytes": output.stdout.len(),
        "stderr_bytes": output.stderr.len(),
        "stdout_tail": &output.stdout,
        "stderr_tail": &output.stderr,
        "stdout_truncated": output.stdout_truncated,
        "stderr_truncated": output.stderr_truncated,
        "log_backpressured": output.log_backpressured,
        "staged_artifact": output.staged_artifact,
        "large_bytes_uploaded": large_bytes_uploaded,
        "registration_response": registration_response,
        "heartbeat_response": heartbeat_response,
        "capability_response": capability_response,
        "task_assignment_response": task_assignment_response,
        "debug_command_response": debug_command_response,
        "log_event_response": log_event_response,
        "vfs_metadata_response": vfs_metadata_response,
        "session_requests": session_requests,
        "coordinator_response": coordinator_response,
    })
}

fn cancelled_node_report(
    _args: &Args,
    task: &RuntimeTask,
    registration_response: Value,
    heartbeat_response: Value,
    capability_response: Value,
    task_assignment_response: Value,
    debug_command_response: Value,
    coordinator_response: Value,
    session_requests: usize,
) -> Value {
    json!({
        "node_status": "cancelled",
        "virtual_thread": &task.task,
        "terminal_state": "cancelled",
        "status_code": null,
        "stdout_bytes": 0,
        "stderr_bytes": 0,
        "stdout_tail": "",
        "stderr_tail": "",
        "stdout_truncated": false,
        "stderr_truncated": false,
        "log_backpressured": false,
        "staged_artifact": null,
        "large_bytes_uploaded": false,
        "registration_response": registration_response,
        "heartbeat_response": heartbeat_response,
        "capability_response": capability_response,
        "task_assignment_response": task_assignment_response,
        "debug_command_response": debug_command_response,
        "log_event_response": null,
        "vfs_metadata_response": null,
        "session_requests": session_requests,
        "coordinator_response": coordinator_response,
    })
}

fn register_node(
    session: &mut CoordinatorSession,
    args: &Args,
) -> Result<Value, Box<dyn std::error::Error>> {
    let public_key = args
        .public_key
        .clone()
        .unwrap_or_else(|| format!("{}-public-key", args.node));
    if let Some(grant) = &args.enrollment_grant {
        session.request(json!({
            "type": "exchange_node_enrollment_grant",
            "tenant": &args.tenant,
            "project": &args.project,
            "node": &args.node,
            "public_key": public_key,
            "enrollment_grant": grant,
            "now_epoch_seconds": 0,
        }))
    } else {
        session.request(json!({
            "type": "attach_node",
            "tenant": &args.tenant,
            "project": &args.project,
            "node": &args.node,
            "public_key": public_key,
        }))
    }
}

fn wait_for_cancellation(
    session: &mut CoordinatorSession,
    args: &Args,
    task: &RuntimeTask,
) -> Result<bool, Box<dyn std::error::Error>> {
    let deadline = Instant::now() + Duration::from_millis(args.control_poll_ms);
    loop {
        let control = session.request(json!({
            "type": "poll_task_control",
            "tenant": &args.tenant,
            "project": &args.project,
            "process": &task.process,
            "node": &args.node,
            "task": &task.task,
        }))?;
        if control
            .get("cancel_requested")
            .and_then(Value::as_bool)
            .unwrap_or(false)
        {
            return Ok(true);
        }
        let now = Instant::now();
        if now >= deadline {
            return Ok(false);
        }
        std::thread::sleep((deadline - now).min(Duration::from_millis(100)));
    }
}

struct CoordinatorSession {
    writer: TcpStream,
    reader: BufReader<TcpStream>,
    requests: usize,
}

impl CoordinatorSession {
    fn connect(addr: &str) -> Result<Self, Box<dyn std::error::Error>> {
        let transport_addr = json_line_transport_addr(addr);
        let writer = TcpStream::connect(&transport_addr)?;
        let reader = BufReader::new(writer.try_clone()?);
        Ok(Self {
            writer,
            reader,
            requests: 0,
        })
    }

    fn request(&mut self, value: Value) -> Result<Value, Box<dyn std::error::Error>> {
        serde_json::to_writer(&mut self.writer, &value)?;
        self.writer.write_all(b"\n")?;
        self.writer.flush()?;

        let mut line = String::new();
        if self.reader.read_line(&mut line)? == 0 {
            return Err("coordinator closed session without a response".into());
        }
        let response: Value = serde_json::from_str(&line)?;
        self.requests += 1;
        if response.get("type").and_then(Value::as_str) == Some("error") {
            return Err(format!("coordinator error: {response}").into());
        }
        Ok(response)
    }

    fn requests(&self) -> usize {
        self.requests
    }
}

fn json_line_transport_addr(endpoint: &str) -> String {
    let endpoint = endpoint.trim();
    for (scheme, default_port) in [("https://", 443), ("http://", 80)] {
        if let Some(rest) = endpoint.strip_prefix(scheme) {
            let authority = rest.split('/').next().unwrap_or(rest);
            if authority.contains(':') {
                return authority.to_owned();
            }
            return format!("{authority}:{default_port}");
        }
    }
    endpoint.to_owned()
}

fn parse_args() -> Result<Args, Box<dyn std::error::Error>> {
    let mut coordinator = None;
    let mut tenant = "tenant".to_owned();
    let mut project = "project".to_owned();
    let mut node = "node".to_owned();
    let mut process = "process".to_owned();
    let mut task = "compile-linux".to_owned();
    let mut command = None;
    let mut command_args = Vec::new();
    let mut artifact = "/vfs/artifacts/node-output.txt".to_owned();
    let mut enrollment_grant = None;
    let mut public_key = None;
    let mut control_poll_ms = 0;
    let mut assignment_poll_ms = 500;
    let mut emit_ready = false;
    let mut worker = false;

    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--coordinator" => coordinator = args.next(),
            "--tenant" => tenant = args.next().ok_or("--tenant requires a value")?,
            "--project-id" => project = args.next().ok_or("--project-id requires a value")?,
            "--node" => node = args.next().ok_or("--node requires a value")?,
            "--process" => process = args.next().ok_or("--process requires a value")?,
            "--task" => task = args.next().ok_or("--task requires a value")?,
            "--command" => command = args.next(),
            "--arg" => command_args.push(args.next().ok_or("--arg requires a value")?),
            "--artifact" => artifact = args.next().ok_or("--artifact requires a value")?,
            "--enrollment-grant" => enrollment_grant = args.next(),
            "--public-key" => public_key = args.next(),
            "--control-poll-ms" => {
                control_poll_ms = args
                    .next()
                    .ok_or("--control-poll-ms requires a value")?
                    .parse()?
            }
            "--assignment-poll-ms" => {
                assignment_poll_ms = args
                    .next()
                    .ok_or("--assignment-poll-ms requires a value")?
                    .parse()?
            }
            "--emit-ready" => emit_ready = true,
            "--worker" => worker = true,
            "--project" => {
                let project_path = PathBuf::from(args.next().ok_or("--project requires a path")?);
                command_args.extend([
                    "test".to_owned(),
                    "--manifest-path".to_owned(),
                    project_path
                        .join("Cargo.toml")
                        .to_string_lossy()
                        .into_owned(),
                ]);
            }
            other => return Err(format!("unknown argument: {other}").into()),
        }
    }

    Ok(Args {
        coordinator: coordinator.ok_or("--coordinator is required")?,
        tenant,
        project,
        node,
        process,
        task,
        command: command.unwrap_or_else(|| "cargo".to_owned()),
        command_args,
        artifact,
        enrollment_grant,
        public_key,
        control_poll_ms,
        assignment_poll_ms,
        emit_ready,
        worker,
    })
}

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

    #[test]
    fn hosted_operator_url_maps_to_json_line_transport_address() {
        assert_eq!(
            json_line_transport_addr("https://disasmer.michelpaulissen.com:9443"),
            "disasmer.michelpaulissen.com:9443"
        );
        assert_eq!(
            json_line_transport_addr("https://disasmer.michelpaulissen.com:9443/auth/device"),
            "disasmer.michelpaulissen.com:9443"
        );
        assert_eq!(
            json_line_transport_addr("http://operator.example.test"),
            "operator.example.test:80"
        );
        assert_eq!(json_line_transport_addr("127.0.0.1:7999"), "127.0.0.1:7999");
    }
}