Publish Clusterflux 1d0b6fa filtered source

crates/clusterflux-coordinator/src/service/logs.rs

@@ -0,0 +1,669 @@
+use clusterflux_core::{
+    ArtifactFlush, ArtifactId, Digest, NodeId, ProcessId, ProjectId, TaskBoundaryValue,
+    TaskInstanceId, TaskJoinResult, TaskJoinState, TenantId, UserId, VfsPath,
+};
+
+use crate::CoordinatorError;
+
+use super::keys::{process_control_key, task_control_key, task_restart_key};
+use super::protocol::TaskAttemptState;
+use super::{
+    artifact_id_from_path, CoordinatorResponse, CoordinatorService, CoordinatorServiceError,
+    TaskCompletionEvent, TaskTerminalState, MAX_TASK_LOG_TAIL_BYTES,
+};
+
+impl CoordinatorService {
+    pub(super) fn handle_report_task_log(
+        &mut self,
+        tenant: String,
+        project: String,
+        process: String,
+        node: String,
+        task: String,
+        stdout_bytes: u64,
+        stderr_bytes: u64,
+        stdout_tail: String,
+        stderr_tail: String,
+        stdout_truncated: bool,
+        stderr_truncated: bool,
+        backpressured: bool,
+    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
+        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)?;
+        validate_task_log_tail("stdout_tail", &stdout_tail)?;
+        validate_task_log_tail("stderr_tail", &stderr_tail)?;
+        let reported_bytes = checked_reported_log_bytes(stdout_bytes, stderr_bytes)?;
+        let now_epoch_seconds = self.current_epoch_seconds()?;
+        self.quota
+            .can_charge_log_bytes(&tenant, &project, reported_bytes, now_epoch_seconds)?;
+        self.quota
+            .charge_log_bytes(&tenant, &project, reported_bytes, now_epoch_seconds)?;
+        Ok(CoordinatorResponse::TaskLogRecorded {
+            process,
+            task,
+            stdout_bytes,
+            stderr_bytes,
+            stdout_tail: if stdout_truncated {
+                format!("{stdout_tail}\n... truncated")
+            } else {
+                stdout_tail
+            },
+            stderr_tail: if stderr_truncated {
+                format!("{stderr_tail}\n... truncated")
+            } else {
+                stderr_tail
+            },
+            backpressured,
+        })
+    }
+
+    pub(super) fn handle_report_vfs_metadata(
+        &mut self,
+        tenant: String,
+        project: String,
+        process: String,
+        node: String,
+        task: String,
+        artifact_path: Option<String>,
+        artifact_digest: Option<Digest>,
+        artifact_size_bytes: Option<u64>,
+        large_bytes_uploaded: bool,
+    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
+        let artifact_path = artifact_path
+            .map(VfsPath::new)
+            .transpose()
+            .map_err(|err| CoordinatorServiceError::InvalidArtifactPath(err.to_string()))?;
+        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)?;
+        if let (Some(path), Some(digest)) = (&artifact_path, artifact_digest) {
+            self.flush_artifact_metadata(ArtifactFlush {
+                id: artifact_id_from_path(path),
+                tenant,
+                project,
+                process: process.clone(),
+                producer_task: task.clone(),
+                retaining_node: node,
+                digest,
+                size: artifact_size_bytes.unwrap_or_default(),
+            })?;
+        }
+        Ok(CoordinatorResponse::VfsMetadataRecorded {
+            process,
+            task,
+            artifact_path,
+            large_bytes_uploaded,
+        })
+    }
+
+    pub(super) fn handle_task_completed(
+        &mut self,
+        tenant: String,
+        project: String,
+        process: String,
+        node: String,
+        task: String,
+        terminal_state: Option<TaskTerminalState>,
+        status_code: Option<i32>,
+        stdout_bytes: u64,
+        stderr_bytes: u64,
+        stdout_tail: String,
+        stderr_tail: String,
+        stdout_truncated: bool,
+        stderr_truncated: bool,
+        artifact_path: Option<String>,
+        artifact_digest: Option<Digest>,
+        artifact_size_bytes: Option<u64>,
+        result: Option<TaskBoundaryValue>,
+    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
+        validate_task_log_tail("stdout_tail", &stdout_tail)?;
+        validate_task_log_tail("stderr_tail", &stderr_tail)?;
+        let artifact_path = artifact_path
+            .map(VfsPath::new)
+            .transpose()
+            .map_err(|err| CoordinatorServiceError::InvalidArtifactPath(err.to_string()))?;
+        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 checkpoint = self
+            .task_restart_checkpoints
+            .get(&super::keys::task_restart_key(
+                &tenant, &project, &process, &task,
+            ))
+            .ok_or_else(|| {
+                CoordinatorError::Unauthorized(
+                    "signed node task completion does not name a coordinator-issued task instance"
+                        .to_owned(),
+                )
+            })?;
+        if checkpoint.assignment.node != node {
+            return Err(CoordinatorError::Unauthorized(
+                "signed node task completion came from a node other than the assigned node"
+                    .to_owned(),
+            )
+            .into());
+        }
+        let mut event = TaskCompletionEvent {
+            tenant,
+            project,
+            process,
+            node,
+            executor: super::TaskExecutor::Node,
+            task_definition: checkpoint.assignment.task_spec.task_definition.clone(),
+            task,
+            attempt_id: None,
+            placement: None,
+            terminal_state: terminal_state
+                .unwrap_or_else(|| TaskTerminalState::from_status_code(status_code)),
+            status_code,
+            stdout_bytes,
+            stderr_bytes,
+            stdout_tail,
+            stderr_tail,
+            stdout_truncated,
+            stderr_truncated,
+            artifact_path,
+            artifact_digest: artifact_digest.clone(),
+            artifact_size_bytes,
+            result,
+        };
+        let reported_bytes = checked_reported_log_bytes(event.stdout_bytes, event.stderr_bytes)?;
+        let now_epoch_seconds = self.current_epoch_seconds()?;
+        self.quota.can_charge_log_bytes(
+            &event.tenant,
+            &event.project,
+            reported_bytes,
+            now_epoch_seconds,
+        )?;
+        self.quota.charge_log_bytes(
+            &event.tenant,
+            &event.project,
+            reported_bytes,
+            now_epoch_seconds,
+        )?;
+        let task_key = task_control_key(
+            &event.tenant,
+            &event.project,
+            &event.process,
+            &event.node,
+            &event.task,
+        );
+        let process_key = process_control_key(&event.tenant, &event.project, &event.process);
+        let process_was_aborted = self.process_aborts.contains(&process_key);
+        event.placement = self.task_placements.remove(&task_key);
+        if let (Some(path), Some(digest)) = (&event.artifact_path, artifact_digest) {
+            self.flush_artifact_metadata(ArtifactFlush {
+                id: artifact_id_from_path(path),
+                tenant: event.tenant.clone(),
+                project: event.project.clone(),
+                process: event.process.clone(),
+                producer_task: event.task.clone(),
+                retaining_node: event.node.clone(),
+                digest,
+                size: artifact_size_bytes.unwrap_or(stdout_bytes),
+            })?;
+        }
+        self.task_cancellations.remove(&task_key);
+        self.task_aborts.remove(&task_key);
+        self.debug_commands.remove(&task_key);
+        self.active_tasks.remove(&task_key);
+        let no_active_tasks =
+            !self
+                .active_tasks
+                .iter()
+                .any(|(task_tenant, task_project, task_process, _, _)| {
+                    task_tenant == &event.tenant
+                        && task_project == &event.project
+                        && task_process == &event.process
+                });
+        let completed_after_main = matches!(event.terminal_state, TaskTerminalState::Completed)
+            && self.task_events.iter().rev().any(|retained| {
+                retained.tenant == event.tenant
+                    && retained.project == event.project
+                    && retained.process == event.process
+                    && matches!(retained.executor, super::TaskExecutor::CoordinatorMain)
+                    && matches!(retained.terminal_state, TaskTerminalState::Completed)
+            });
+        if no_active_tasks {
+            self.process_aborts.remove(&process_key);
+            if (self.process_cancellations.remove(&process_key) || completed_after_main)
+                && !self.main_runtime.controls.contains_key(&process_key)
+            {
+                self.coordinator
+                    .abort_process(&event.tenant, &event.project, &event.process)?;
+                self.clear_debug_state_for_process(&event.tenant, &event.project, &event.process);
+                self.clear_operator_panel_state(&event.tenant, &event.project, &event.process);
+            }
+        }
+        if process_was_aborted {
+            let checkpoint_key = super::keys::task_restart_key(
+                &event.tenant,
+                &event.project,
+                &event.process,
+                &event.task,
+            );
+            self.task_restart_checkpoints.remove(&checkpoint_key);
+            self.task_restart_checkpoint_order
+                .retain(|retained| retained != &checkpoint_key);
+        }
+        let awaiting_operator = self.finish_task_attempt(&mut event);
+        self.record_task_completion_event(event.clone());
+        if !awaiting_operator {
+            self.notify_coordinator_main_waiters(&event);
+        }
+        Ok(CoordinatorResponse::TaskRecorded {
+            process: event.process,
+            task: event.task,
+            events_recorded: self.task_events.len(),
+        })
+    }
+
+    pub(super) fn handle_list_task_events(
+        &mut self,
+        tenant: String,
+        project: String,
+        actor_user: String,
+        process: Option<String>,
+    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
+        let tenant = TenantId::new(tenant);
+        let project = ProjectId::new(project);
+        let _actor = UserId::new(actor_user);
+        let process = process.map(ProcessId::new);
+        if let Some(process) = &process {
+            self.authorize_task_event_process_scope(&tenant, &project, process)?;
+        }
+        let events = self
+            .task_events
+            .iter()
+            .filter(|event| {
+                event.tenant == tenant
+                    && event.project == project
+                    && process
+                        .as_ref()
+                        .is_none_or(|process| event.process == *process)
+            })
+            .cloned()
+            .collect();
+        Ok(CoordinatorResponse::TaskEvents { events })
+    }
+
+    pub(super) fn handle_list_task_snapshots(
+        &mut self,
+        tenant: String,
+        project: String,
+        actor_user: String,
+        process: String,
+    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
+        let tenant = TenantId::new(tenant);
+        let project = ProjectId::new(project);
+        let _actor = UserId::new(actor_user);
+        let process = ProcessId::new(process);
+        self.authorize_task_event_process_scope(&tenant, &project, &process)?;
+        let snapshots = self
+            .task_attempts
+            .iter()
+            .filter(
+                |((attempt_tenant, attempt_project, attempt_process, _), _)| {
+                    attempt_tenant == &tenant
+                        && attempt_project == &project
+                        && attempt_process == &process
+                },
+            )
+            .flat_map(|(_, attempts)| attempts.iter().cloned())
+            .collect();
+        Ok(CoordinatorResponse::TaskSnapshots { snapshots })
+    }
+
+    fn authorize_task_event_process_scope(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        process: &ProcessId,
+    ) -> Result<(), CoordinatorServiceError> {
+        let active_in_scope = self
+            .coordinator
+            .active_process(tenant, project, process)
+            .is_some();
+        let historical_in_scope = self.task_events.iter().any(|event| {
+            event.tenant == *tenant && event.project == *project && event.process == *process
+        }) || self.process_scope_history.iter().any(
+            |(historical_tenant, historical_project, historical_process)| {
+                historical_tenant == tenant
+                    && historical_project == project
+                    && historical_process == process
+            },
+        );
+        let process_exists_outside_scope = self
+            .coordinator
+            .active_process_exists_outside_scope(tenant, project, process)
+            || self.task_events.iter().any(|event| {
+                event.process == *process && (event.tenant != *tenant || event.project != *project)
+            })
+            || self.process_scope_history.iter().any(
+                |(historical_tenant, historical_project, historical_process)| {
+                    historical_process == process
+                        && (historical_tenant != tenant || historical_project != project)
+                },
+            );
+        if !active_in_scope && !historical_in_scope && process_exists_outside_scope {
+            return Err(CoordinatorError::Unauthorized(
+                "task event access is outside the virtual process tenant/project scope".to_owned(),
+            )
+            .into());
+        }
+        Ok(())
+    }
+
+    pub(super) fn handle_join_task(
+        &mut self,
+        tenant: String,
+        project: String,
+        actor_user: String,
+        process: String,
+        task: String,
+    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
+        let tenant = TenantId::new(tenant);
+        let project = ProjectId::new(project);
+        let _actor = UserId::new(actor_user);
+        let process = ProcessId::new(process);
+        let task = TaskInstanceId::new(task);
+        Ok(CoordinatorResponse::TaskJoined {
+            join: self.task_join_result(tenant, project, process, task),
+        })
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(super) fn handle_join_child_task(
+        &mut self,
+        tenant: String,
+        project: String,
+        process: String,
+        node: String,
+        parent_task: String,
+        task: String,
+    ) -> Result<CoordinatorResponse, CoordinatorServiceError> {
+        let tenant = TenantId::new(tenant);
+        let project = ProjectId::new(project);
+        let process = ProcessId::new(process);
+        let node = NodeId::new(node);
+        let parent_task = TaskInstanceId::new(parent_task);
+        self.authorize_node_for_process_or_termination(&node, &tenant, &project, &process)?;
+        if !self.active_tasks.contains(&super::keys::task_control_key(
+            &tenant,
+            &project,
+            &process,
+            &node,
+            &parent_task,
+        )) {
+            return Err(CoordinatorError::Unauthorized(
+                "child task join requires a currently active parent task on the signed node"
+                    .to_owned(),
+            )
+            .into());
+        }
+        Ok(CoordinatorResponse::TaskJoined {
+            join: self.task_join_result(tenant, project, process, TaskInstanceId::new(task)),
+        })
+    }
+
+    pub(super) fn task_join_result(
+        &self,
+        tenant: TenantId,
+        project: ProjectId,
+        process: ProcessId,
+        task: TaskInstanceId,
+    ) -> TaskJoinResult {
+        let attempt_key = task_restart_key(&tenant, &project, &process, &task);
+        if self
+            .task_attempts
+            .get(&attempt_key)
+            .is_some_and(|attempts| {
+                attempts
+                    .iter()
+                    .rev()
+                    .find(|attempt| attempt.current)
+                    .is_some_and(|attempt| {
+                        matches!(
+                            attempt.state,
+                            TaskAttemptState::Queued
+                                | TaskAttemptState::Running
+                                | TaskAttemptState::FailedAwaitingAction
+                        )
+                    })
+            })
+        {
+            return TaskJoinResult::pending(
+                process,
+                task,
+                "logical task is still running or awaiting operator action",
+            );
+        }
+        let event = self.task_events.iter().rev().find(|event| {
+            event.tenant == tenant
+                && event.project == project
+                && event.process == process
+                && event.task == task
+        });
+
+        if let Some(event) = event {
+            TaskJoinResult::from_remote_completion(
+                event.process.clone(),
+                event.task.clone(),
+                event.node.clone(),
+                join_state_for_terminal(&event.terminal_state),
+                event.result.clone(),
+                event.status_code,
+                join_message_for_event(event),
+            )
+        } else {
+            let known = self.task_is_known_or_active(&tenant, &project, &process, &task);
+            TaskJoinResult::pending(
+                process,
+                task,
+                if known {
+                    "waiting for signed node task_completed event before join returns"
+                } else {
+                    "no signed node completion event has been observed for this task"
+                },
+            )
+        }
+    }
+
+    pub(super) fn record_task_completion_event(&mut self, mut event: TaskCompletionEvent) {
+        event.stdout_tail = bounded_log_tail(event.stdout_tail, &mut event.stdout_truncated);
+        event.stderr_tail = bounded_log_tail(event.stderr_tail, &mut event.stderr_truncated);
+        let process_scope = (
+            event.tenant.clone(),
+            event.project.clone(),
+            event.process.clone(),
+        );
+        self.process_scope_history
+            .retain(|retained| retained != &process_scope);
+        while self.process_scope_history.len() >= super::MAX_TASK_EVENTS_TOTAL {
+            self.process_scope_history.pop_front();
+        }
+        self.process_scope_history.push_back(process_scope);
+        while self
+            .task_events
+            .iter()
+            .filter(|retained| {
+                retained.tenant == event.tenant
+                    && retained.project == event.project
+                    && retained.process == event.process
+            })
+            .count()
+            >= super::MAX_TASK_EVENTS_PER_PROCESS
+        {
+            let Some(index) = self.task_events.iter().position(|retained| {
+                retained.tenant == event.tenant
+                    && retained.project == event.project
+                    && retained.process == event.process
+            }) else {
+                break;
+            };
+            self.task_events.remove(index);
+        }
+        while self.task_events.len() >= super::MAX_TASK_EVENTS_TOTAL {
+            self.task_events.pop_front();
+        }
+        self.task_events.push_back(event);
+    }
+
+    fn finish_task_attempt(&mut self, event: &mut TaskCompletionEvent) -> bool {
+        let key = task_restart_key(&event.tenant, &event.project, &event.process, &event.task);
+        let Some(attempt) = self
+            .task_attempts
+            .get_mut(&key)
+            .and_then(|attempts| attempts.iter_mut().rev().find(|attempt| attempt.current))
+        else {
+            return false;
+        };
+        event.attempt_id = Some(attempt.attempt_id.clone());
+        attempt.status_code = event.status_code;
+        attempt.artifact_path = event.artifact_path.clone();
+        attempt.artifact_digest = event.artifact_digest.clone();
+        attempt.artifact_size_bytes = event.artifact_size_bytes;
+        attempt.error = (!event.stderr_tail.trim().is_empty()).then(|| event.stderr_tail.clone());
+        let awaiting_operator = event.terminal_state == TaskTerminalState::Failed
+            && attempt.failure_policy == clusterflux_core::TaskFailurePolicy::AwaitOperator;
+        attempt.state = if awaiting_operator {
+            TaskAttemptState::FailedAwaitingAction
+        } else {
+            match event.terminal_state {
+                TaskTerminalState::Completed => TaskAttemptState::Completed,
+                TaskTerminalState::Failed => TaskAttemptState::Failed,
+                TaskTerminalState::Cancelled => TaskAttemptState::Cancelled,
+            }
+        };
+        attempt.command_state = Some(if awaiting_operator {
+            "failed_awaiting_action".to_owned()
+        } else {
+            format!("{:?}", event.terminal_state).to_ascii_lowercase()
+        });
+        awaiting_operator
+    }
+
+    fn flush_artifact_metadata(
+        &mut self,
+        flush: ArtifactFlush,
+    ) -> Result<(), CoordinatorServiceError> {
+        let now_epoch_seconds = self.current_epoch_seconds()?;
+        self.artifact_registry
+            .expire_download_links(now_epoch_seconds);
+        let pinned = self
+            .task_restart_checkpoints
+            .values()
+            .flat_map(|checkpoint| checkpoint.assignment.task_spec.required_artifacts.iter())
+            .chain(
+                self.pending_task_launches
+                    .iter()
+                    .flat_map(|pending| pending.task_spec.required_artifacts.iter()),
+            )
+            .cloned()
+            .collect::<std::collections::BTreeSet<ArtifactId>>();
+        self.artifact_registry
+            .flush_metadata_bounded(flush, &pinned)
+            .map(|_| ())
+            .map_err(CoordinatorServiceError::Protocol)
+    }
+
+    fn task_is_known_or_active(
+        &self,
+        tenant: &TenantId,
+        project: &ProjectId,
+        process: &ProcessId,
+        task: &TaskInstanceId,
+    ) -> bool {
+        self.active_tasks
+            .iter()
+            .any(|(task_tenant, task_project, task_process, _, task_id)| {
+                task_tenant == tenant
+                    && task_project == project
+                    && task_process == process
+                    && task_id == task
+            })
+            || self.pending_task_launches.iter().any(|pending| {
+                &pending.tenant == tenant
+                    && &pending.project == project
+                    && &pending.process == process
+                    && &pending.task == task
+            })
+            || self.task_assignments.values().any(|assignments| {
+                assignments.iter().any(|assignment| {
+                    &assignment.tenant == tenant
+                        && &assignment.project == project
+                        && &assignment.process == process
+                        && &assignment.task == task
+                })
+            })
+    }
+}
+
+fn checked_reported_log_bytes(
+    stdout_bytes: u64,
+    stderr_bytes: u64,
+) -> Result<u64, CoordinatorServiceError> {
+    stdout_bytes.checked_add(stderr_bytes).ok_or_else(|| {
+        CoordinatorServiceError::Protocol(
+            "reported task log byte counts exceed the supported range".to_owned(),
+        )
+    })
+}
+
+fn validate_task_log_tail(kind: &str, value: &str) -> Result<(), CoordinatorServiceError> {
+    if value.len() > MAX_TASK_LOG_TAIL_BYTES {
+        return Err(CoordinatorServiceError::InvalidTaskLogTail(format!(
+            "{kind} is {} bytes; max is {MAX_TASK_LOG_TAIL_BYTES}",
+            value.len()
+        )));
+    }
+    Ok(())
+}
+
+fn bounded_log_tail(mut value: String, truncated: &mut bool) -> String {
+    if value.len() <= MAX_TASK_LOG_TAIL_BYTES {
+        return value;
+    }
+    let mut boundary = MAX_TASK_LOG_TAIL_BYTES;
+    while !value.is_char_boundary(boundary) {
+        boundary -= 1;
+    }
+    value.truncate(boundary);
+    *truncated = true;
+    value
+}
+
+fn join_state_for_terminal(terminal: &TaskTerminalState) -> TaskJoinState {
+    match terminal {
+        TaskTerminalState::Completed => TaskJoinState::Completed,
+        TaskTerminalState::Failed => TaskJoinState::Failed,
+        TaskTerminalState::Cancelled => TaskJoinState::Cancelled,
+    }
+}
+
+fn join_message_for_event(event: &TaskCompletionEvent) -> String {
+    match event.terminal_state {
+        TaskTerminalState::Completed => {
+            "joined result from signed node task_completed event".to_owned()
+        }
+        TaskTerminalState::Failed => {
+            let stderr = event.stderr_tail.trim();
+            if stderr.is_empty() {
+                "remote task failed".to_owned()
+            } else {
+                format!("remote task failed: {stderr}")
+            }
+        }
+        TaskTerminalState::Cancelled => "remote task was cancelled".to_owned(),
+    }
+}