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

use std::collections::{BTreeMap, BTreeSet};
use std::io::Write;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use clusterflux_core::{
    sign_node_request, signed_request_payload_digest, ArtifactId, Digest, NodeCapabilities, NodeId,
    TaskSpec, MIN_SIGNED_NODE_POLL_INTERVAL_MS,
};
use serde_json::{json, Value};

use crate::assignment_runner::run_verified_wasmtime_assignment;
#[cfg(test)]
use crate::coordinator_session::control_endpoint_identity;
use crate::coordinator_session::CoordinatorSession;
use crate::debug_agent::poll_task_cancellation;
use crate::node_identity::{
    establish_node_identity, node_nonce, node_private_key_for_runtime, signed_node_request_json,
    unix_timestamp_seconds,
};
#[cfg(test)]
use crate::node_identity::{load_or_create_local_node_credential, unix_timestamp_nanos};
use crate::source_snapshot::snapshot_project;
use crate::task_artifacts::{
    clean_stale_task_output_roots, current_epoch_seconds, NodeArtifactRetentionLimits,
    NodeArtifactStore,
};
use crate::task_reports::{record_cancelled_task, record_completed_task, record_failed_task};

static WORKER_SHUTDOWN_REQUESTED: AtomicBool = AtomicBool::new(false);

pub(crate) fn worker_shutdown_requested() -> bool {
    WORKER_SHUTDOWN_REQUESTED.load(Ordering::Acquire)
}

#[cfg(unix)]
extern "C" fn request_worker_shutdown(_signal: libc::c_int) {
    WORKER_SHUTDOWN_REQUESTED.store(true, Ordering::SeqCst);
}

fn install_worker_shutdown_handler() {
    WORKER_SHUTDOWN_REQUESTED.store(false, Ordering::SeqCst);
    #[cfg(unix)]
    unsafe {
        libc::signal(libc::SIGINT, request_worker_shutdown as libc::sighandler_t);
        libc::signal(libc::SIGTERM, request_worker_shutdown as libc::sighandler_t);
    }
}

#[derive(Clone, Debug)]
pub(crate) struct Args {
    pub(crate) coordinator: String,
    pub(crate) tenant: String,
    pub(crate) project: String,
    pub(crate) project_root: Option<PathBuf>,
    pub(crate) node: String,
    pub(crate) enrollment_grant: Option<String>,
    pub(crate) public_key: Option<String>,
    pub(crate) control_poll_ms: u64,
    pub(crate) assignment_poll_ms: u64,
    pub(crate) emit_ready: bool,
    pub(crate) worker: bool,
}

#[derive(Clone, Debug)]
pub(crate) struct RuntimeTask {
    pub(crate) process: String,
    pub(crate) task: String,
    pub(crate) epoch: Option<u64>,
    pub(crate) task_spec: Option<TaskSpec>,
    pub(crate) bundle_digest: Option<Digest>,
    pub(crate) wasm_module_base64: Option<String>,
    pub(crate) task_assignment_response: Value,
}

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

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

    let node_private_key = node_private_key_for_runtime(&args.node)?;
    let registration = establish_node_identity(&mut session, &args, &node_private_key)?;
    let heartbeat_request = json!({
        "type": "node_heartbeat",
        "node": &args.node,
    });
    let heartbeat_signature = sign_node_request(
        &node_private_key,
        &NodeId::from(args.node.as_str()),
        "node_heartbeat",
        &signed_request_payload_digest(&heartbeat_request),
        node_nonce("node-heartbeat"),
        unix_timestamp_seconds(),
    )
    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
    let mut heartbeat_request = heartbeat_request;
    heartbeat_request["node_signature"] = json!(heartbeat_signature);
    let heartbeat = session.request(heartbeat_request)?;
    clean_stale_task_output_roots(args.project_root.as_deref(), &args.node)?;
    let artifact_store = NodeArtifactStore::for_runtime(args.project_root.as_deref(), &args.node)?;
    let retention_limits = NodeArtifactRetentionLimits::from_environment()?;
    artifact_store.garbage_collect(retention_limits, &BTreeSet::new(), current_epoch_seconds())?;
    let capability_report = report_node_capabilities(
        &args,
        &mut session,
        &node_private_key,
        &artifact_store,
        true,
    )?;

    if !args.worker {
        return Err(
            "one-shot native command mode was removed; run `clusterflux-node --worker` and launch a bundled Wasm task through the coordinator"
                .into(),
        );
    }
    install_worker_shutdown_handler();
    worker_loop(
        &args,
        &mut session,
        registration,
        heartbeat,
        capability_report,
        &node_private_key,
    )
}

fn worker_loop(
    args: &Args,
    session: &mut CoordinatorSession,
    registration: Value,
    heartbeat: Value,
    capability_report: Value,
    node_private_key: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    const ARTIFACT_GC_INTERVAL: Duration = Duration::from_secs(30);
    let artifact_store = NodeArtifactStore::for_runtime(args.project_root.as_deref(), &args.node)?;
    let retention_limits = NodeArtifactRetentionLimits::from_environment()?;
    let mut last_artifact_gc = Instant::now();
    let mut restart_pins = BTreeMap::<ArtifactId, Instant>::new();
    if args.emit_ready {
        println!(
            "{}",
            serde_json::to_string(&json!({
                "node_status": "ready",
                "mode": "worker",
                "node": &args.node,
            }))?
        );
        std::io::stdout().flush()?;
    }

    loop {
        if worker_shutdown_requested() {
            report_node_capabilities(args, session, node_private_key, &artifact_store, false)?;
            return Ok(());
        }
        if service_pending_artifact_transfer(args, session, node_private_key)? {
            continue;
        }
        if last_artifact_gc.elapsed() >= ARTIFACT_GC_INTERVAL {
            let now = Instant::now();
            restart_pins.retain(|_, expiry| *expiry > now);
            let pinned = restart_pins.keys().cloned().collect::<BTreeSet<_>>();
            artifact_store.garbage_collect(retention_limits, &pinned, current_epoch_seconds())?;
            report_node_capabilities(args, session, node_private_key, &artifact_store, true)?;
            last_artifact_gc = Instant::now();
        }
        let response = session.request(signed_node_request_json(
            args,
            node_private_key,
            "poll_task_assignment",
            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)?;
        if args.emit_ready {
            println!(
                "{}",
                serde_json::to_string(&json!({
                    "node_status": "assignment_started",
                    "node": &args.node,
                    "process": &runtime_task.process,
                    "virtual_thread": &runtime_task.task,
                    "task_assignment_response": &runtime_task.task_assignment_response,
                }))?
            );
            std::io::stdout().flush()?;
        }
        if let Some(task_spec) = &runtime_task.task_spec {
            let expiry = Instant::now()
                .checked_add(Duration::from_secs(retention_limits.restart_pin_seconds))
                .unwrap_or_else(Instant::now);
            for artifact in &task_spec.required_artifacts {
                restart_pins.insert(artifact.clone(), expiry);
            }
        }
        let report = run_runtime_task(
            args,
            session,
            runtime_task,
            registration.clone(),
            heartbeat.clone(),
            capability_report.clone(),
            node_private_key,
        )?;
        if let Some(artifact) = report
            .pointer("/vfs_metadata_response/artifact_path")
            .and_then(Value::as_str)
            .and_then(|path| path.strip_prefix("/vfs/artifacts/"))
        {
            let expiry = Instant::now()
                .checked_add(Duration::from_secs(retention_limits.restart_pin_seconds))
                .unwrap_or_else(Instant::now);
            restart_pins.insert(ArtifactId::new(artifact), expiry);
        }
        println!("{}", serde_json::to_string(&report)?);
        std::io::stdout().flush()?;
    }
}

fn report_node_capabilities(
    args: &Args,
    session: &mut CoordinatorSession,
    node_private_key: &str,
    artifact_store: &NodeArtifactStore,
    online: bool,
) -> Result<Value, Box<dyn std::error::Error>> {
    let artifact_locations = artifact_store
        .artifact_ids()?
        .into_iter()
        .map(|artifact| artifact.as_str().to_owned())
        .collect::<Vec<_>>();
    let source_snapshots = args
        .project_root
        .as_deref()
        .map(snapshot_project)
        .transpose()?
        .into_iter()
        .map(|snapshot| snapshot.digest)
        .collect::<Vec<_>>();
    session.request(signed_node_request_json(
        args,
        node_private_key,
        "report_node_capabilities",
        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": source_snapshots,
            "artifact_locations": artifact_locations,
            "direct_connectivity": false,
            "online": online,
        }),
    )?)
}

fn service_pending_artifact_transfer(
    args: &Args,
    session: &mut CoordinatorSession,
    node_private_key: &str,
) -> Result<bool, Box<dyn std::error::Error>> {
    let response = session.request(signed_node_request_json(
        args,
        node_private_key,
        "poll_artifact_transfer",
        json!({
            "type": "poll_artifact_transfer",
            "tenant": &args.tenant,
            "project": &args.project,
            "node": &args.node,
        }),
    )?)?;
    let Some(transfer) = response.get("transfer").filter(|value| !value.is_null()) else {
        return Ok(false);
    };
    let transfer_id = required_string(transfer, "transfer_id")?;
    let artifact = ArtifactId::new(required_string(transfer, "artifact")?);
    let expected_digest: Digest = serde_json::from_value(
        transfer
            .get("expected_digest")
            .cloned()
            .ok_or("artifact transfer omitted expected_digest")?,
    )?;
    let expected_size_bytes = transfer
        .get("expected_size_bytes")
        .and_then(Value::as_u64)
        .ok_or("artifact transfer omitted expected_size_bytes")?;
    let offset = transfer
        .get("offset")
        .and_then(Value::as_u64)
        .ok_or("artifact transfer omitted offset")?;
    let max_chunk_bytes = transfer
        .get("max_chunk_bytes")
        .and_then(Value::as_u64)
        .ok_or("artifact transfer omitted max_chunk_bytes")?;
    let read = NodeArtifactStore::for_runtime(args.project_root.as_deref(), &args.node).and_then(
        |store| {
            store.read_verified_chunk(
                &artifact,
                &expected_digest,
                expected_size_bytes,
                offset,
                max_chunk_bytes,
            )
        },
    );
    let content = match read {
        Ok(content) => content,
        Err(message) => {
            session.request(signed_node_request_json(
                args,
                node_private_key,
                "fail_artifact_transfer",
                json!({
                    "type": "fail_artifact_transfer",
                    "tenant": &args.tenant,
                    "project": &args.project,
                    "node": &args.node,
                    "transfer_id": transfer_id,
                    "artifact": artifact,
                    "message": message,
                }),
            )?)?;
            return Ok(true);
        }
    };
    let eof = offset.saturating_add(content.len() as u64) == expected_size_bytes;
    session.request(signed_node_request_json(
        args,
        node_private_key,
        "upload_artifact_transfer_chunk",
        json!({
            "type": "upload_artifact_transfer_chunk",
            "tenant": &args.tenant,
            "project": &args.project,
            "node": &args.node,
            "transfer_id": transfer_id,
            "artifact": artifact,
            "offset": offset,
            "content_base64": BASE64_STANDARD.encode(&content),
            "chunk_digest": Digest::sha256(&content),
            "eof": eof,
        }),
    )?)?;
    Ok(true)
}

pub(crate) fn runtime_task_from_assignment(
    value: &Value,
) -> Result<RuntimeTask, Box<dyn std::error::Error>> {
    let task_spec: TaskSpec = serde_json::from_value(
        value
            .get("task_spec")
            .cloned()
            .ok_or("task assignment missing task_spec")?,
    )?;
    Ok(RuntimeTask {
        process: required_string(value, "process")?,
        task: required_string(value, "task")?,
        epoch: value.get("epoch").and_then(Value::as_u64),
        bundle_digest: task_spec.bundle_digest.clone(),
        task_spec: Some(task_spec),
        wasm_module_base64: Some(required_string(value, "wasm_module_base64")?),
        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,
    node_private_key: &str,
) -> 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(signed_node_request_json(
        args,
        node_private_key,
        "reconnect_node",
        json!({
        "type": "reconnect_node",
        "node": &args.node,
        "process": &task.process,
        "epoch": epoch,
        }),
    )?)?;
    let debug_command = session.request(signed_node_request_json(
        args,
        node_private_key,
        "poll_debug_command",
        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 && poll_task_cancellation(session, args, &task, node_private_key)? {
        return record_cancelled_task(
            args,
            session,
            &task,
            registration,
            heartbeat,
            capability_report,
            debug_command,
            node_private_key,
        );
    }

    let execution = run_verified_wasmtime_assignment(args, &task, node_private_key);
    match execution {
        Ok((output, manifest, result)) => match crate::task_artifacts::retained_result_artifact(
            args.project_root.as_deref(),
            &args.node,
            result.as_ref(),
        ) {
            Ok(retained) => record_completed_task(
                args,
                session,
                task,
                output,
                manifest,
                result,
                retained,
                registration,
                heartbeat,
                capability_report,
                debug_command,
                node_private_key,
            ),
            Err(error) => record_failed_task(
                args,
                session,
                &task,
                registration,
                heartbeat,
                capability_report,
                debug_command,
                node_private_key,
                &error,
            ),
        },
        Err(error) => {
            let error = error.to_string();
            if error.contains("task execution cancelled:") {
                record_cancelled_task(
                    args,
                    session,
                    &task,
                    registration,
                    heartbeat,
                    capability_report,
                    debug_command,
                    node_private_key,
                )
            } else {
                record_failed_task(
                    args,
                    session,
                    &task,
                    registration,
                    heartbeat,
                    capability_report,
                    debug_command,
                    node_private_key,
                    &error,
                )
            }
        }
    }
}

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 project_root = None;
    let mut node = "node".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")?,
            "--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 = validate_assignment_poll_ms(
                    args.next()
                        .ok_or("--assignment-poll-ms requires a value")?
                        .parse()?,
                )?
            }
            "--emit-ready" => emit_ready = true,
            "--worker" => worker = true,
            "--project-root" => {
                project_root = Some(PathBuf::from(
                    args.next().ok_or("--project-root requires a path")?,
                ));
            }
            other => return Err(format!("unknown argument: {other}").into()),
        }
    }

    Ok(Args {
        coordinator: coordinator.ok_or("--coordinator is required")?,
        tenant,
        project,
        project_root,
        node,
        enrollment_grant,
        public_key,
        control_poll_ms,
        assignment_poll_ms,
        emit_ready,
        worker,
    })
}

fn validate_assignment_poll_ms(value: u64) -> Result<u64, String> {
    if value < MIN_SIGNED_NODE_POLL_INTERVAL_MS {
        return Err(format!(
            "--assignment-poll-ms must be at least {MIN_SIGNED_NODE_POLL_INTERVAL_MS} ms so signed polling remains within the coordinator's bounded replay window"
        ));
    }
    Ok(value)
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeSet;

    use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
    use clusterflux_core::{
        Digest, ProcessId, ProjectId, TaskBoundaryValue, TaskDispatch, TaskInstanceId, TaskSpec,
        TenantId, WasmExportAbi, WasmTaskResult,
    };

    use crate::assignment_runner::run_verified_wasmtime_assignment;
    use clusterflux_core::MIN_SIGNED_NODE_POLL_INTERVAL_MS;

    use super::{
        control_endpoint_identity, load_or_create_local_node_credential,
        validate_assignment_poll_ms, Args, RuntimeTask,
    };

    #[test]
    fn node_poll_interval_cannot_exhaust_the_bounded_replay_window() {
        assert!(validate_assignment_poll_ms(MIN_SIGNED_NODE_POLL_INTERVAL_MS).is_ok());
        let error = validate_assignment_poll_ms(MIN_SIGNED_NODE_POLL_INTERVAL_MS - 1).unwrap_err();
        assert!(error.contains("bounded replay window"));
    }

    #[test]
    fn hosted_url_remains_an_https_control_endpoint() {
        assert_eq!(
            control_endpoint_identity("https://clusterflux.michelpaulissen.com").unwrap(),
            "https://clusterflux.michelpaulissen.com/api/v1/control"
        );
        assert_eq!(
            control_endpoint_identity("https://clusterflux.michelpaulissen.com/api/v1/control")
                .unwrap(),
            "https://clusterflux.michelpaulissen.com/api/v1/control"
        );
        assert_eq!(
            control_endpoint_identity("127.0.0.1:7999").unwrap(),
            "clusterflux+tcp://127.0.0.1:7999"
        );
    }

    #[test]
    fn daemon_local_node_credential_is_durable_between_runs() {
        let temp = std::env::temp_dir().join(format!(
            "clusterflux-node-credential-test-{}-{}",
            std::process::id(),
            super::unix_timestamp_nanos()
        ));
        std::fs::create_dir_all(&temp).unwrap();
        let first = load_or_create_local_node_credential(&temp, "daemon-node").unwrap();
        let second = load_or_create_local_node_credential(&temp, "daemon-node").unwrap();

        assert_eq!(first, second);
        assert!(temp.join(".clusterflux").join("nodes").exists());

        std::fs::remove_dir_all(&temp).unwrap();
    }

    #[test]
    fn daemon_wasm_task_assignment_uses_abi_version_and_verifies_bundle_digest() {
        let args = test_args();
        let task_instance = TaskInstanceId::from("task_add_one-1");
        let boundary = TaskBoundaryValue::SmallJson(serde_json::json!(2));
        let result = serde_json::to_string(&WasmTaskResult::completed(
            task_instance.clone(),
            boundary.clone(),
        ))
        .unwrap();
        let wat_result = result.replace('\\', "\\\\").replace('"', "\\\"");
        let result_length = result.len();
        let packed = ((result_length as u64) << 32) | 2048;
        let wasm = wat::parse_str(format!(
            r#"(module
                  (memory (export "memory") 1)
                  (data (i32.const 2048) "{wat_result}")
                  (func (export "clusterflux_alloc_v1") (param i32) (result i32)
                    i32.const 1024)
                  (func (export "task_add_one") (param i32 i32) (result i64)
                    i64.const {packed}))"#
        ))
        .unwrap();
        let task = RuntimeTask {
            process: "vp".to_owned(),
            task: task_instance.as_str().to_owned(),
            epoch: Some(7),
            task_spec: Some(TaskSpec {
                tenant: TenantId::from("tenant"),
                project: ProjectId::from("project"),
                process: ProcessId::from("vp"),
                task_definition: clusterflux_core::TaskDefinitionId::from("task_add_one"),
                task_instance,
                dispatch: TaskDispatch::CoordinatorNodeWasm {
                    export: Some("task_add_one".to_owned()),
                    abi: WasmExportAbi::TaskV1,
                },
                environment_id: None,
                environment: None,
                environment_digest: None,
                required_capabilities: BTreeSet::new(),
                dependency_cache: None,
                source_snapshot: None,
                required_artifacts: Vec::new(),
                args: Vec::new(),
                vfs_epoch: 7,
                failure_policy: Default::default(),
                bundle_digest: Some(Digest::sha256(&wasm)),
            }),
            bundle_digest: Some(Digest::sha256(&wasm)),
            wasm_module_base64: Some(BASE64_STANDARD.encode(&wasm)),
            task_assignment_response: serde_json::json!({}),
        };

        let (output, manifest, result) =
            run_verified_wasmtime_assignment(&args, &task, "test-node-private-key").unwrap();

        assert_eq!(output.status_code, Some(0));
        assert_eq!(
            output.stdout,
            format!("{}\n", serde_json::to_string(&boundary).unwrap())
        );
        assert!(output.staged_artifact.is_none());
        assert!(manifest.objects.is_empty());
        assert!(!manifest.large_bytes_uploaded);
        assert_eq!(result, Some(boundary));

        let mismatch = RuntimeTask {
            bundle_digest: Some(Digest::sha256("different bundle bytes")),
            wasm_module_base64: Some(BASE64_STANDARD.encode("not valid wasm")),
            ..task
        };
        let error = run_verified_wasmtime_assignment(&args, &mismatch, "test-node-private-key")
            .unwrap_err();
        assert!(error.to_string().contains("bundle digest mismatch"));
        assert!(!error.to_string().contains("failed to parse"));
    }

    fn test_args() -> Args {
        Args {
            coordinator: "127.0.0.1:1".to_owned(),
            tenant: "tenant".to_owned(),
            project: "project".to_owned(),
            project_root: None,
            node: "node".to_owned(),
            enrollment_grant: None,
            public_key: None,
            control_poll_ms: 0,
            assignment_poll_ms: 1,
            emit_ready: false,
            worker: false,
        }
    }
}