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, ProcessId, ProjectId, TaskInstanceId, TaskSpec, TenantId, MIN_SIGNED_NODE_POLL_INTERVAL_MS, }; use serde_json::{json, Value}; use crate::assignment_runner::{assignment_error_log_bytes, 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, pub(crate) node: String, pub(crate) enrollment_grant: Option, pub(crate) public_key: Option, 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, pub(crate) task_spec: Option, pub(crate) bundle_digest: Option, pub(crate) wasm_module_base64: Option, pub(crate) task_assignment_response: Value, } pub(crate) fn run() -> Result<(), Box> { 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", "tenant": &args.tenant, "project": &args.project, "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> { 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::::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::>(); 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::try_new(artifact.to_owned())?, 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> { let artifact_locations = artifact_store .artifact_ids()? .into_iter() .map(|artifact| artifact.as_str().to_owned()) .collect::>(); let source_snapshots = args .project_root .as_deref() .map(snapshot_project) .transpose()? .into_iter() .map(|snapshot| snapshot.digest) .collect::>(); 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> { 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::try_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> { let task_spec: TaskSpec = serde_json::from_value( value .get("task_spec") .cloned() .ok_or("task assignment missing task_spec")?, )?; let process = ProcessId::try_new(required_string(value, "process")?)?; let task = TaskInstanceId::try_new(required_string(value, "task")?)?; Ok(RuntimeTask { process: process.to_string(), task: task.to_string(), 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> { 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> { 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", "tenant": &args.tenant, "project": &args.project, "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, 0, 0, ); } 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, output.stdout_source_bytes, output.stderr_source_bytes, ), }, Err(error) => { let (stdout_source_bytes, stderr_source_bytes) = assignment_error_log_bytes(error.as_ref()); 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, stdout_source_bytes, stderr_source_bytes, ) } else { record_failed_task( args, session, &task, registration, heartbeat, capability_report, debug_command, node_private_key, &error, stdout_source_bytes, stderr_source_bytes, ) } } } } fn parse_args() -> Result> { 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()), } } TenantId::try_new(tenant.clone())?; ProjectId::try_new(project.clone())?; NodeId::try_new(node.clone())?; if let Some(grant) = enrollment_grant.as_ref() { clusterflux_core::validate_opaque_token(grant, 512)?; } 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 { 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.lesstuff.com").unwrap(), "https://clusterflux.lesstuff.com/api/v1/control" ); assert_eq!( control_endpoint_identity("https://clusterflux.lesstuff.com/api/v1/control").unwrap(), "https://clusterflux.lesstuff.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, } } }