use std::collections::{BTreeSet, VecDeque}; use std::time::{SystemTime, UNIX_EPOCH}; use clusterflux_core::{ AgentId, AgentSignedRequest, CredentialKind, DefaultScheduler, Digest, NodeId, PlacementRequest, ProcessId, ProjectId, RendezvousRequest, Scheduler, SourcePreparation, TaskCheckpoint, TaskInstanceId, TaskSpec, TenantId, UserId, }; use crate::CoordinatorError; use super::keys::{process_control_key, task_control_key}; use super::{ CoordinatorResponse, CoordinatorService, CoordinatorServiceError, SourcePreparationDisposition, SourcePreparationStatus, TaskAssignment, TaskCancellationTarget, VirtualProcessStatus, WorkflowActor, }; #[derive(Clone, Debug)] pub(super) struct PendingTaskLaunch { pub(super) tenant: TenantId, pub(super) project: ProjectId, pub(super) process: ProcessId, pub(super) task: TaskInstanceId, pub(super) request: PlacementRequest, pub(super) epoch: u64, pub(super) artifact_path: String, pub(super) task_spec: TaskSpec, pub(super) wasm_module_base64: String, } #[derive(Clone, Debug)] pub(super) struct TaskRestartCheckpoint { pub(super) checkpoint: TaskCheckpoint, pub(super) assignment: TaskAssignment, } impl CoordinatorService { pub(super) fn handle_poll_task_assignment( &mut self, tenant: String, project: String, node: String, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let node = NodeId::new(node); let identity = self .coordinator .node_identity(&node) .ok_or(CoordinatorError::UnknownNode)?; if identity.tenant != tenant || identity.project != project { return Err(CoordinatorError::Unauthorized( "task assignment poll is outside the enrolled tenant/project scope".to_owned(), ) .into()); } let assignment_key = (tenant.clone(), project.clone(), node.clone()); let assignment = self .task_assignments .get_mut(&assignment_key) .and_then(VecDeque::pop_front); if assignment.is_some() { return Ok(CoordinatorResponse::TaskAssignment { assignment: assignment.map(Box::new), }); } let assignment = self.assign_pending_task_to_node(&tenant, &project, &node)?; Ok(CoordinatorResponse::TaskAssignment { assignment: assignment.map(Box::new), }) } fn assign_pending_task_to_node( &mut self, tenant: &TenantId, project: &ProjectId, node: &NodeId, ) -> Result, CoordinatorServiceError> { let Some(descriptor) = self.node_descriptors.get(node).cloned() else { return Ok(None); }; let mut remaining = VecDeque::new(); let mut selected = None; while let Some(pending) = self.pending_task_launches.pop_front() { if selected.is_some() { remaining.push_back(pending); continue; } if &pending.tenant != tenant || &pending.project != project { remaining.push_back(pending); continue; } if self.process_cancellations.contains(&process_control_key( &pending.tenant, &pending.project, &pending.process, )) { continue; } let Some(active) = self.coordinator.active_process( &pending.tenant, &pending.project, &pending.process, ) else { continue; }; if active.tenant != pending.tenant || active.project != pending.project || active.coordinator_epoch != pending.epoch { continue; } let Ok(placement) = DefaultScheduler.place(std::slice::from_ref(&descriptor), &pending.request) else { remaining.push_back(pending); continue; }; let assignment = TaskAssignment { tenant: pending.tenant.clone(), project: pending.project.clone(), process: pending.process.clone(), task: pending.task.clone(), node: placement.node.clone(), epoch: pending.epoch, artifact_path: pending.artifact_path, task_spec: pending.task_spec, wasm_module_base64: pending.wasm_module_base64, }; self.assign_task_attempt(&assignment.task_spec, assignment.node.clone()); self.capture_task_restart_checkpoint(&assignment)?; let task_key = task_control_key( &pending.tenant, &pending.project, &pending.process, &placement.node, &pending.task, ); self.task_placements.insert(task_key.clone(), placement); self.active_tasks.insert(task_key); selected = Some(assignment); } self.pending_task_launches = remaining; Ok(selected) } pub(super) fn handle_request_rendezvous( &mut self, scope: clusterflux_core::DataPlaneScope, source: clusterflux_core::NodeEndpoint, destination: clusterflux_core::NodeEndpoint, direct_connectivity: bool, failure_reason: String, ) -> Result { let now_epoch_seconds = self.current_epoch_seconds()?; let charged_rendezvous_attempts = self.quota.charge_rendezvous_attempt( &scope.tenant, &scope.project, now_epoch_seconds, )?; let plan = self.transport.plan_authenticated_direct_bulk_transfer( RendezvousRequest { scope, source, destination, }, direct_connectivity, failure_reason, )?; Ok(CoordinatorResponse::RendezvousPlan { plan, charged_rendezvous_attempts, }) } pub(super) fn handle_request_source_preparation( &mut self, tenant: String, project: String, provider: clusterflux_core::SourceProviderKind, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let preparation = SourcePreparation::node_task(tenant.clone(), project.clone(), provider); let request = PlacementRequest { tenant, project, environment: None, environment_digest: None, environment_cache_required: false, required_capabilities: preparation.required_capabilities.clone(), dependency_cache: None, source_snapshot: None, required_artifacts: Default::default(), quota_available: true, policy_allowed: true, prefer_node: None, }; let nodes = self.live_node_descriptors(); let disposition = match DefaultScheduler.place(&nodes, &request) { Ok(placement) => SourcePreparationDisposition::Assigned { node: placement.node, }, Err(err) => SourcePreparationDisposition::Pending { reason: if err.message.is_empty() { "waiting for any capable node to prepare source".to_owned() } else { err.message }, }, }; Ok(CoordinatorResponse::SourcePreparation { status: SourcePreparationStatus { preparation, disposition, }, }) } pub(super) fn handle_complete_source_preparation( &mut self, tenant: String, project: String, node: String, provider: clusterflux_core::SourceProviderKind, source_snapshot: Digest, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let node = NodeId::new(node); let identity = self .coordinator .node_identity(&node) .ok_or(CoordinatorError::UnknownNode)?; if identity.tenant != tenant || identity.project != project { return Err(CoordinatorError::Unauthorized( "source preparation completion is outside the enrolled tenant/project scope" .to_owned(), ) .into()); } let descriptor = self.node_descriptors.get_mut(&node).ok_or_else(|| { CoordinatorError::Unauthorized( "source preparation completion requires a node capability report".to_owned(), ) })?; if !descriptor.source_snapshots.contains(&source_snapshot) && descriptor.source_snapshots.len() >= super::MAX_NODE_REPORTED_OBJECTS_PER_KIND { return Err(CoordinatorServiceError::Protocol(format!( "node source snapshot retention limit of {} reached; refresh the node capability report", super::MAX_NODE_REPORTED_OBJECTS_PER_KIND ))); } descriptor.source_snapshots.insert(source_snapshot.clone()); Ok(CoordinatorResponse::SourcePreparationCompleted { node, provider, source_snapshot, }) } pub(super) fn handle_start_process( &mut self, tenant: String, project: String, actor_user: Option, actor_agent: Option, agent_public_key_fingerprint: Option, agent_signature: Option, request_payload_digest: Option<&Digest>, process: String, launch_attempt: Option, restart: bool, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let process = ProcessId::new(process); self.coordinator.ensure_tenant_active(&tenant)?; let actor = self.workflow_actor( &tenant, &project, actor_user, actor_agent, agent_public_key_fingerprint, agent_signature, request_payload_digest, "start_process", &process, None, )?; let replacing_existing = if let Some(active) = self .coordinator .active_process_for_project(&tenant, &project) { if active.id != process || !restart { return Err(CoordinatorError::Unauthorized(format!( "project already has active virtual process {}; attach to or restart it, request cooperative cancellation, abort it, or use another Coordinator Project", active.id )) .into()); } true } else { false }; if !replacing_existing { self.quota.ensure_process_admission( &tenant, self.coordinator.active_process_count_for_tenant(&tenant), )?; } let now_epoch_seconds = self.current_epoch_seconds()?; let charged_spawns = self.quota .charge_workflow_spawn(&tenant, &project, now_epoch_seconds)?; if replacing_existing { self.main_runtime.interrupt_process( &tenant, &project, &process, "virtual process incarnation replaced", ); self.main_runtime .controls .remove(&process_control_key(&tenant, &project, &process)); } self.process_cancellations .remove(&process_control_key(&tenant, &project, &process)); self.process_aborts .remove(&process_control_key(&tenant, &project, &process)); self.clear_debug_state_for_process(&tenant, &project, &process); self.clear_operator_panel_state(&tenant, &project, &process); self.task_cancellations .retain(|(task_tenant, task_project, task_process, _, _)| { task_tenant != &tenant || task_project != &project || task_process != &process }); self.task_aborts .retain(|(task_tenant, task_project, task_process, _, _)| { task_tenant != &tenant || task_project != &project || task_process != &process }); self.active_tasks .retain(|(task_tenant, task_project, task_process, _, _)| { task_tenant != &tenant || task_project != &project || task_process != &process }); self.task_placements .retain(|(task_tenant, task_project, task_process, _, _), _| { task_tenant != &tenant || task_project != &project || task_process != &process }); self.task_assignments.retain(|_, assignments| { assignments.retain(|assignment| { assignment.tenant != tenant || assignment.project != project || assignment.process != process }); !assignments.is_empty() }); self.pending_task_launches.retain(|pending| { pending.tenant != tenant || pending.project != project || pending.process != process }); self.task_restart_checkpoints.retain( |(checkpoint_tenant, checkpoint_project, checkpoint_process, _), _| { checkpoint_tenant != &tenant || checkpoint_project != &project || checkpoint_process != &process }, ); self.task_restart_checkpoint_order.retain( |(checkpoint_tenant, checkpoint_project, checkpoint_process, _)| { checkpoint_tenant != &tenant || checkpoint_project != &project || checkpoint_process != &process }, ); self.task_events.retain(|event| { event.tenant != tenant || event.project != project || event.process != process }); self.task_attempts .retain(|(attempt_tenant, attempt_project, attempt_process, _), _| { attempt_tenant != &tenant || attempt_project != &project || attempt_process != &process }); self.restart_launches .retain(|(attempt_tenant, attempt_project, attempt_process, _)| { attempt_tenant != &tenant || attempt_project != &project || attempt_process != &process }); self.debug_audit_events.retain(|event| { event.tenant != tenant || event.project != project || event.process != process }); let active = self.coordinator.start_process_for_launch_attempt( tenant, project, process.clone(), launch_attempt.map(clusterflux_core::LaunchAttemptId::new), ); Ok(CoordinatorResponse::ProcessStarted { process, launch_attempt: active .launch_attempt .as_ref() .map(|attempt| attempt.as_str().to_owned()), epoch: self.coordinator.coordinator_epoch(), actor, charged_spawns, }) } pub(super) fn handle_reconnect_node( &mut self, node: String, process: String, epoch: u64, ) -> Result { let node = NodeId::new(node); let process = ProcessId::new(process); self.coordinator .reconnect_node(&node, Some((&process, epoch)))?; Ok(CoordinatorResponse::NodeReconnected { node, process }) } pub(super) fn handle_cancel_task( &mut self, tenant: String, project: String, process: String, node: String, task: String, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let process = ProcessId::new(process); let node = NodeId::new(node); let task = TaskInstanceId::new(task); self.coordinator .authorize_node_for_process(&node, &tenant, &project, &process)?; let active = self .coordinator .active_process(&tenant, &project, &process) .ok_or_else(|| { CoordinatorError::Unauthorized( "task cancellation requires an active virtual process".to_owned(), ) })?; if !active.connected_nodes.contains(&node) { return Err(CoordinatorError::Unauthorized( "task cancellation target node is not connected to the virtual process".to_owned(), ) .into()); } self.task_cancellations .insert(task_control_key(&tenant, &project, &process, &node, &task)); Ok(CoordinatorResponse::TaskCancellationRequested { process, task, node, }) } pub(super) fn handle_cancel_process( &mut self, tenant: String, project: String, actor_user: String, process: String, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let process = ProcessId::new(process); let _actor_user = actor_user; let active = self .coordinator .active_process(&tenant, &project, &process) .ok_or_else(|| { CoordinatorError::Unauthorized( "process cancellation requires an active virtual process".to_owned(), ) })?; debug_assert_eq!(active.tenant, tenant); debug_assert_eq!(active.project, project); self.process_cancellations .insert(process_control_key(&tenant, &project, &process)); self.main_runtime.interrupt_process( &tenant, &project, &process, "virtual process cancellation requested", ); self.clear_debug_state_for_process(&tenant, &project, &process); self.pending_task_launches.retain(|pending| { pending.tenant != tenant || pending.project != project || pending.process != process }); let mut cancelled_tasks = Vec::new(); let mut affected_nodes = BTreeSet::new(); for (task_tenant, task_project, task_process, node, task) in self.active_tasks.iter() { if task_tenant == &tenant && task_project == &project && task_process == &process { self.task_cancellations .insert(task_control_key(&tenant, &project, &process, node, task)); affected_nodes.insert(node.clone()); cancelled_tasks.push(TaskCancellationTarget { process: process.clone(), task: task.clone(), node: node.clone(), }); } } let process_key = process_control_key(&tenant, &project, &process); if cancelled_tasks.is_empty() && !self.main_runtime.controls.contains_key(&process_key) { self.coordinator .abort_process(&tenant, &project, &process)?; self.clear_operator_panel_state(&tenant, &project, &process); } Ok(CoordinatorResponse::ProcessCancellationRequested { process, cancelled_tasks, affected_nodes: affected_nodes.into_iter().collect(), }) } pub(super) fn handle_abort_process( &mut self, tenant: String, project: String, actor_user: String, process: String, launch_attempt: Option, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let process = ProcessId::new(process); let _actor_user = actor_user; let active = self .coordinator .active_process(&tenant, &project, &process) .ok_or_else(|| { CoordinatorError::Unauthorized( "process abort requires an active virtual process".to_owned(), ) })?; debug_assert_eq!(active.tenant, tenant); debug_assert_eq!(active.project, project); let launch_attempt = launch_attempt.map(clusterflux_core::LaunchAttemptId::new); if let Some(expected) = launch_attempt.as_ref() { if active.launch_attempt.as_ref() != Some(expected) { return Err(CoordinatorError::Unauthorized(format!( "launch rollback denied: attempt {} does not own process {}", expected.as_str(), process.as_str() )) .into()); } } let process_key = process_control_key(&tenant, &project, &process); self.process_cancellations.remove(&process_key); self.task_cancellations .retain(|(task_tenant, task_project, task_process, _, _)| { task_tenant != &tenant || task_project != &project || task_process != &process }); self.process_aborts.insert(process_key); self.main_runtime .interrupt_process(&tenant, &project, &process, "virtual process aborted"); self.main_runtime .controls .remove(&process_control_key(&tenant, &project, &process)); self.clear_debug_state_for_process(&tenant, &project, &process); self.clear_operator_panel_state(&tenant, &project, &process); self.pending_task_launches.retain(|pending| { pending.tenant != tenant || pending.project != project || pending.process != process }); self.task_assignments.retain(|_, assignments| { assignments.retain(|assignment| { assignment.tenant != tenant || assignment.project != project || assignment.process != process }); !assignments.is_empty() }); let mut aborted_tasks = Vec::new(); let mut affected_nodes = BTreeSet::new(); for (task_tenant, task_project, task_process, node, task) in self.active_tasks.iter() { if task_tenant == &tenant && task_project == &project && task_process == &process { self.task_aborts .insert(task_control_key(&tenant, &project, &process, node, task)); affected_nodes.insert(node.clone()); aborted_tasks.push(TaskCancellationTarget { process: process.clone(), task: task.clone(), node: node.clone(), }); } } if let Some(launch_attempt) = launch_attempt.as_ref() { self.coordinator.abort_process_for_launch_attempt( &tenant, &project, &process, launch_attempt, )?; } else { self.coordinator .abort_process(&tenant, &project, &process)?; } let active_restart_tasks = aborted_tasks .iter() .map(|target| target.task.clone()) .collect::>(); self.task_restart_checkpoints.retain( |(checkpoint_tenant, checkpoint_project, checkpoint_process, checkpoint_task), _| { checkpoint_tenant != &tenant || checkpoint_project != &project || checkpoint_process != &process || active_restart_tasks.contains(checkpoint_task) }, ); self.task_restart_checkpoint_order.retain( |(checkpoint_tenant, checkpoint_project, checkpoint_process, checkpoint_task)| { checkpoint_tenant != &tenant || checkpoint_project != &project || checkpoint_process != &process || active_restart_tasks.contains(checkpoint_task) }, ); if aborted_tasks.is_empty() { self.process_aborts .remove(&process_control_key(&tenant, &project, &process)); } Ok(CoordinatorResponse::ProcessAborted { process, aborted_tasks, affected_nodes: affected_nodes.into_iter().collect(), }) } pub(super) fn handle_list_processes( &mut self, tenant: String, project: String, actor_user: String, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let actor = UserId::new(actor_user); let processes = self .coordinator .active_processes_for_project(&tenant, &project) .into_iter() .map(|active| { let process_key = process_control_key(&active.tenant, &active.project, &active.id); let main = self.main_runtime.controls.get(&process_key); let state = if self.process_cancellations.contains(&process_key) { "cancelling" } else { main.map_or("running", |main| main.state.as_str()) }; let main_wait_state = main.and_then(|main| { if main.state != "running" { return None; } if self.pending_task_launches.iter().any(|pending| { pending.tenant == active.tenant && pending.project == active.project && pending.process == active.id }) { Some("waiting_for_node".to_owned()) } else if self.main_runtime.is_waiting_for_task( &active.tenant, &active.project, &active.id, ) { Some("waiting_for_task".to_owned()) } else { Some("executing".to_owned()) } }); VirtualProcessStatus { process: active.id, state: state.to_owned(), main_task_definition: main.map(|main| main.task_definition.clone()), main_task_instance: main.map(|main| main.task_instance.clone()), main_state: main.map(|main| main.state.clone()), main_wait_state, main_debug_epoch: main.and_then(|main| main.debug.requested_epoch()), connected_nodes: active.connected_nodes.into_iter().collect(), coordinator_epoch: active.coordinator_epoch, } }) .collect(); Ok(CoordinatorResponse::ProcessStatuses { processes, actor }) } pub(super) fn handle_poll_task_control( &mut self, tenant: String, project: String, process: String, node: String, task: String, ) -> Result { let tenant = TenantId::new(tenant); let project = ProjectId::new(project); let process = ProcessId::new(process); let node = NodeId::new(node); let task = TaskInstanceId::new(task); self.authorize_node_for_process_or_termination(&node, &tenant, &project, &process)?; let cancel_requested = self .task_cancellations .contains(&task_control_key(&tenant, &project, &process, &node, &task)) || self .process_cancellations .contains(&process_control_key(&tenant, &project, &process)); let abort_requested = self .task_aborts .contains(&task_control_key(&tenant, &project, &process, &node, &task)) || self .process_aborts .contains(&process_control_key(&tenant, &project, &process)); Ok(CoordinatorResponse::TaskControl { process, task, cancel_requested, abort_requested, }) } pub(super) fn workflow_actor( &mut self, tenant: &TenantId, project: &ProjectId, actor_user: Option, actor_agent: Option, agent_public_key_fingerprint: Option, agent_signature: Option, request_payload_digest: Option<&Digest>, request_kind: &str, process: &ProcessId, task: Option<&TaskInstanceId>, ) -> Result { if let Some(agent) = actor_agent { let agent = AgentId::new(agent); let signature = agent_signature.ok_or_else(|| { CoordinatorError::Unauthorized( "agent workflow dispatch requires a signed request proving private-key possession" .to_owned(), ) })?; let request_payload_digest = request_payload_digest.ok_or_else(|| { CoordinatorError::Unauthorized( "agent workflow dispatch requires a canonical signed request payload" .to_owned(), ) })?; if signature.nonce.trim().is_empty() || signature.nonce.len() > 256 { return Err(CoordinatorError::Unauthorized( "agent signed request nonce is missing or invalid".to_owned(), ) .into()); } let now_epoch_seconds = unix_timestamp_seconds(); let replay_key = ( tenant.clone(), project.clone(), agent.clone(), signature.nonce.clone(), ); const AGENT_SIGNATURE_WINDOW_SECONDS: u64 = 300; self.agent_replay_nonces.retain(|_, issued_at| { now_epoch_seconds <= issued_at.saturating_add(AGENT_SIGNATURE_WINDOW_SECONDS) }); if self.agent_replay_nonces.contains_key(&replay_key) { return Err(CoordinatorError::Unauthorized( "agent signed request nonce has already been used".to_owned(), ) .into()); } let canonical_scope = clusterflux_core::AgentWorkflowRequestScope::new( tenant.clone(), project.clone(), request_kind, process.clone(), task.cloned(), ) .map_err(CoordinatorError::Unauthorized)?; let record = self.coordinator.authorize_agent_project_run( canonical_scope.for_agent(&agent), agent_public_key_fingerprint.as_ref(), request_payload_digest, &signature, now_epoch_seconds, )?; if self .agent_replay_nonces .keys() .filter(|(retained_tenant, retained_project, retained_agent, _)| { retained_tenant == tenant && retained_project == project && retained_agent == &agent }) .count() >= super::MAX_REPLAY_NONCES_PER_AUTHORITY { return Err(CoordinatorError::Unauthorized( "agent signed request replay window is full; retry after the bounded signature window advances" .to_owned(), ) .into()); } self.agent_replay_nonces .insert(replay_key, signature.issued_at_epoch_seconds); return Ok(WorkflowActor { kind: "agent".to_owned(), user: Some(record.user), agent: Some(agent), credential_kind: CredentialKind::PublicKey, public_key_fingerprint: Some(record.public_key_fingerprint), authenticated_without_browser: true, scopes: record.scopes, }); } let actor = UserId::new(actor_user.unwrap_or_else(|| "user".to_owned())); Ok(WorkflowActor { kind: "user".to_owned(), user: Some(actor), agent: None, credential_kind: CredentialKind::BrowserSession, public_key_fingerprint: None, authenticated_without_browser: false, scopes: vec!["project:read".to_owned(), "project:run".to_owned()], }) } } fn unix_timestamp_seconds() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .map(|duration| duration.as_secs()) .unwrap_or(0) }