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

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(),
    }
}