use crate::task::{TaskId, TaskState};
use crate::workflow::{WorkflowId, WorkflowState};
use chrono::{DateTime, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
fn serialize_atomic<S: Serializer>(val: &Arc<AtomicU64>, ser: S) -> Result<S::Ok, S::Error> {
ser.serialize_u64(val.load(Ordering::Relaxed))
}
fn deserialize_atomic<'de, D: Deserializer<'de>>(de: D) -> Result<Arc<AtomicU64>, D::Error> {
let n = u64::deserialize(de)?;
Ok(Arc::new(AtomicU64::new(n)))
}
fn default_atomic() -> Arc<AtomicU64> {
Arc::new(AtomicU64::new(0))
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskMetrics {
pub task_id: TaskId,
pub task_name: String,
pub state: TaskState,
pub start_time: Option<DateTime<Utc>>,
pub end_time: Option<DateTime<Utc>>,
pub duration: Option<Duration>,
pub retry_count: u32,
pub error: Option<String>,
}
impl TaskMetrics {
#[must_use]
pub fn new(task_id: TaskId, task_name: String) -> Self {
Self {
task_id,
task_name,
state: TaskState::Pending,
start_time: None,
end_time: None,
duration: None,
retry_count: 0,
error: None,
}
}
pub fn mark_started(&mut self) {
self.state = TaskState::Running;
self.start_time = Some(Utc::now());
}
pub fn mark_completed(&mut self) {
self.state = TaskState::Completed;
self.end_time = Some(Utc::now());
if let Some(start) = self.start_time {
self.duration = Some(Duration::from_millis(
u64::try_from((Utc::now() - start).num_milliseconds()).unwrap_or(0),
));
}
}
pub fn mark_failed(&mut self, error: String) {
self.state = TaskState::Failed;
self.end_time = Some(Utc::now());
self.error = Some(error);
if let Some(start) = self.start_time {
self.duration = Some(Duration::from_millis(
u64::try_from((Utc::now() - start).num_milliseconds()).unwrap_or(0),
));
}
}
pub fn increment_retry(&mut self) {
self.retry_count += 1;
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowMetrics {
pub workflow_id: WorkflowId,
pub workflow_name: String,
pub state: WorkflowState,
pub start_time: Option<DateTime<Utc>>,
pub end_time: Option<DateTime<Utc>>,
pub duration: Option<Duration>,
pub tasks: HashMap<TaskId, TaskMetrics>,
pub total_tasks: usize,
#[serde(
serialize_with = "serialize_atomic",
deserialize_with = "deserialize_atomic",
default = "default_atomic"
)]
pub completed_tasks: Arc<AtomicU64>,
#[serde(
serialize_with = "serialize_atomic",
deserialize_with = "deserialize_atomic",
default = "default_atomic"
)]
pub failed_tasks: Arc<AtomicU64>,
#[serde(
serialize_with = "serialize_atomic",
deserialize_with = "deserialize_atomic",
default = "default_atomic"
)]
pub running_tasks: Arc<AtomicU64>,
}
impl WorkflowMetrics {
#[must_use]
pub fn new(workflow_id: WorkflowId, workflow_name: String, total_tasks: usize) -> Self {
Self {
workflow_id,
workflow_name,
state: WorkflowState::Created,
start_time: None,
end_time: None,
duration: None,
tasks: HashMap::new(),
total_tasks,
completed_tasks: Arc::new(AtomicU64::new(0)),
failed_tasks: Arc::new(AtomicU64::new(0)),
running_tasks: Arc::new(AtomicU64::new(0)),
}
}
#[must_use]
pub fn completed_tasks_count(&self) -> u64 {
self.completed_tasks.load(Ordering::Relaxed)
}
#[must_use]
pub fn failed_tasks_count(&self) -> u64 {
self.failed_tasks.load(Ordering::Relaxed)
}
#[must_use]
pub fn running_tasks_count(&self) -> u64 {
self.running_tasks.load(Ordering::Relaxed)
}
pub fn mark_started(&mut self) {
self.state = WorkflowState::Running;
self.start_time = Some(Utc::now());
}
pub fn mark_completed(&mut self) {
self.state = WorkflowState::Completed;
self.end_time = Some(Utc::now());
if let Some(start) = self.start_time {
self.duration = Some(Duration::from_millis(
u64::try_from((Utc::now() - start).num_milliseconds()).unwrap_or(0),
));
}
}
pub fn mark_failed(&mut self) {
self.state = WorkflowState::Failed;
self.end_time = Some(Utc::now());
if let Some(start) = self.start_time {
self.duration = Some(Duration::from_millis(
u64::try_from((Utc::now() - start).num_milliseconds()).unwrap_or(0),
));
}
}
pub fn update_task(&mut self, task_metrics: TaskMetrics) {
if let Some(old_metrics) = self.tasks.get(&task_metrics.task_id) {
self.update_counters(old_metrics.state, false);
}
self.update_counters(task_metrics.state, true);
self.tasks.insert(task_metrics.task_id, task_metrics);
}
fn update_counters(&self, state: TaskState, increment: bool) {
match state {
TaskState::Completed => {
if increment {
self.completed_tasks.fetch_add(1, Ordering::Relaxed);
} else {
let _ = self.completed_tasks.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|v| if v > 0 { Some(v - 1) } else { Some(0) },
);
}
}
TaskState::Failed => {
if increment {
self.failed_tasks.fetch_add(1, Ordering::Relaxed);
} else {
let _ =
self.failed_tasks
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| {
if v > 0 {
Some(v - 1)
} else {
Some(0)
}
});
}
}
TaskState::Running | TaskState::Retrying => {
if increment {
self.running_tasks.fetch_add(1, Ordering::Relaxed);
} else {
let _ = self.running_tasks.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|v| if v > 0 { Some(v - 1) } else { Some(0) },
);
}
}
_ => {}
}
}
#[must_use]
pub fn progress_percentage(&self) -> f64 {
if self.total_tasks == 0 {
return 100.0;
}
self.completed_tasks.load(Ordering::Relaxed) as f64 / self.total_tasks as f64 * 100.0
}
#[must_use]
pub fn average_task_duration(&self) -> Option<Duration> {
let durations: Vec<_> = self.tasks.values().filter_map(|m| m.duration).collect();
if durations.is_empty() {
return None;
}
let total_ms: u64 = durations.iter().map(|d| d.as_millis() as u64).sum();
Some(Duration::from_millis(total_ms / durations.len() as u64))
}
#[must_use]
pub fn throughput(&self) -> f64 {
let completed = self.completed_tasks.load(Ordering::Relaxed) as f64;
if let (Some(start), Some(end)) = (self.start_time, self.end_time) {
let duration_secs = (end - start).num_seconds().max(1);
completed / duration_secs as f64
} else if let Some(start) = self.start_time {
let duration_secs = (Utc::now() - start).num_seconds().max(1);
completed / duration_secs as f64
} else {
0.0
}
}
}
pub struct MonitoringService {
workflows: Arc<DashMap<WorkflowId, WorkflowMetrics>>,
history: Arc<DashMap<WorkflowId, WorkflowMetrics>>,
}
impl MonitoringService {
#[must_use]
pub fn new() -> Self {
Self {
workflows: Arc::new(DashMap::new()),
history: Arc::new(DashMap::new()),
}
}
pub fn start_workflow(
&self,
workflow_id: WorkflowId,
workflow_name: String,
total_tasks: usize,
) {
let mut metrics = WorkflowMetrics::new(workflow_id, workflow_name, total_tasks);
metrics.mark_started();
self.workflows.insert(workflow_id, metrics);
}
pub fn update_task(
&self,
workflow_id: WorkflowId,
task_id: TaskId,
task_name: String,
state: TaskState,
error: Option<String>,
) {
if let Some(mut workflow_metrics) = self.workflows.get_mut(&workflow_id) {
let mut task_metrics = workflow_metrics
.tasks
.get(&task_id)
.cloned()
.unwrap_or_else(|| TaskMetrics::new(task_id, task_name.clone()));
match state {
TaskState::Running => task_metrics.mark_started(),
TaskState::Completed => task_metrics.mark_completed(),
TaskState::Failed => {
if let Some(err) = error {
task_metrics.mark_failed(err);
} else {
task_metrics.mark_failed("Unknown error".to_string());
}
}
TaskState::Retrying => task_metrics.increment_retry(),
_ => task_metrics.state = state,
}
workflow_metrics.update_task(task_metrics);
}
}
pub fn complete_workflow(&self, workflow_id: WorkflowId, success: bool) {
if let Some((_, mut metrics)) = self.workflows.remove(&workflow_id) {
if success {
metrics.mark_completed();
} else {
metrics.mark_failed();
}
self.history.insert(workflow_id, metrics);
}
}
#[must_use]
pub fn get_workflow_metrics(&self, workflow_id: &WorkflowId) -> Option<WorkflowMetrics> {
self.workflows
.get(workflow_id)
.map(|m| m.clone())
.or_else(|| self.history.get(workflow_id).map(|m| m.clone()))
}
#[must_use]
pub fn get_active_workflows(&self) -> Vec<WorkflowMetrics> {
self.workflows
.iter()
.map(|entry| entry.value().clone())
.collect()
}
#[must_use]
pub fn get_history(&self, limit: Option<usize>) -> Vec<WorkflowMetrics> {
let mut history: Vec<_> = self
.history
.iter()
.map(|entry| entry.value().clone())
.collect();
history.sort_by(|a, b| b.start_time.cmp(&a.start_time));
if let Some(limit) = limit {
history.truncate(limit);
}
history
}
pub fn clear_history(&self, older_than: DateTime<Utc>) {
self.history.retain(|_, metrics| {
if let Some(end_time) = metrics.end_time {
end_time > older_than
} else {
true
}
});
}
#[must_use]
pub fn get_statistics(&self) -> SystemStatistics {
let active_workflows = self.workflows.len();
let total_workflows = self.history.len() + active_workflows;
let completed_workflows = self
.history
.iter()
.filter(|entry| matches!(entry.value().state, WorkflowState::Completed))
.count();
let failed_workflows = self
.history
.iter()
.filter(|entry| matches!(entry.value().state, WorkflowState::Failed))
.count();
let total_tasks_completed: usize = (self
.workflows
.iter()
.map(|entry| entry.value().completed_tasks.load(Ordering::Relaxed))
.sum::<u64>()
+ self
.history
.iter()
.map(|entry| entry.value().completed_tasks.load(Ordering::Relaxed))
.sum::<u64>()) as usize;
let total_tasks_failed: usize = (self
.workflows
.iter()
.map(|entry| entry.value().failed_tasks.load(Ordering::Relaxed))
.sum::<u64>()
+ self
.history
.iter()
.map(|entry| entry.value().failed_tasks.load(Ordering::Relaxed))
.sum::<u64>()) as usize;
SystemStatistics {
active_workflows,
total_workflows,
completed_workflows,
failed_workflows,
total_tasks_completed,
total_tasks_failed,
}
}
}
impl Default for MonitoringService {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemStatistics {
pub active_workflows: usize,
pub total_workflows: usize,
pub completed_workflows: usize,
pub failed_workflows: usize,
pub total_tasks_completed: usize,
pub total_tasks_failed: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_task_metrics_creation() {
let task_id = TaskId::new();
let metrics = TaskMetrics::new(task_id, "test-task".to_string());
assert_eq!(metrics.task_id, task_id);
assert_eq!(metrics.task_name, "test-task");
assert_eq!(metrics.state, TaskState::Pending);
}
#[test]
fn test_task_metrics_lifecycle() {
let task_id = TaskId::new();
let mut metrics = TaskMetrics::new(task_id, "test-task".to_string());
metrics.mark_started();
assert_eq!(metrics.state, TaskState::Running);
assert!(metrics.start_time.is_some());
metrics.mark_completed();
assert_eq!(metrics.state, TaskState::Completed);
assert!(metrics.end_time.is_some());
assert!(metrics.duration.is_some());
}
#[test]
fn test_task_metrics_failure() {
let task_id = TaskId::new();
let mut metrics = TaskMetrics::new(task_id, "test-task".to_string());
metrics.mark_started();
metrics.mark_failed("Test error".to_string());
assert_eq!(metrics.state, TaskState::Failed);
assert_eq!(metrics.error, Some("Test error".to_string()));
}
#[test]
fn test_workflow_metrics_creation() {
let workflow_id = WorkflowId::new();
let metrics = WorkflowMetrics::new(workflow_id, "test-workflow".to_string(), 5);
assert_eq!(metrics.workflow_id, workflow_id);
assert_eq!(metrics.total_tasks, 5);
assert_eq!(metrics.completed_tasks_count(), 0);
}
#[test]
fn test_workflow_metrics_progress() {
let workflow_id = WorkflowId::new();
let metrics = WorkflowMetrics::new(workflow_id, "test-workflow".to_string(), 10);
assert_eq!(metrics.progress_percentage(), 0.0);
metrics.completed_tasks.store(5, Ordering::Relaxed);
assert_eq!(metrics.progress_percentage(), 50.0);
metrics.completed_tasks.store(10, Ordering::Relaxed);
assert_eq!(metrics.progress_percentage(), 100.0);
}
#[test]
fn test_monitoring_service_creation() {
let service = MonitoringService::new();
assert_eq!(service.get_active_workflows().len(), 0);
}
#[test]
fn test_monitoring_service_workflow_tracking() {
let service = MonitoringService::new();
let workflow_id = WorkflowId::new();
service.start_workflow(workflow_id, "test-workflow".to_string(), 3);
let metrics = service.get_workflow_metrics(&workflow_id);
assert!(metrics.is_some());
assert_eq!(
metrics.expect("should succeed in test").state,
WorkflowState::Running
);
}
#[test]
fn test_monitoring_service_task_updates() {
let service = MonitoringService::new();
let workflow_id = WorkflowId::new();
let task_id = TaskId::new();
service.start_workflow(workflow_id, "test-workflow".to_string(), 1);
service.update_task(
workflow_id,
task_id,
"task-1".to_string(),
TaskState::Running,
None,
);
let metrics = service
.get_workflow_metrics(&workflow_id)
.expect("should succeed in test");
assert_eq!(metrics.running_tasks_count(), 1);
service.update_task(
workflow_id,
task_id,
"task-1".to_string(),
TaskState::Completed,
None,
);
let metrics = service
.get_workflow_metrics(&workflow_id)
.expect("should succeed in test");
assert_eq!(metrics.completed_tasks_count(), 1);
assert_eq!(metrics.running_tasks_count(), 0);
}
#[test]
fn test_monitoring_service_completion() {
let service = MonitoringService::new();
let workflow_id = WorkflowId::new();
service.start_workflow(workflow_id, "test-workflow".to_string(), 1);
assert_eq!(service.get_active_workflows().len(), 1);
service.complete_workflow(workflow_id, true);
assert_eq!(service.get_active_workflows().len(), 0);
let history = service.get_history(None);
assert_eq!(history.len(), 1);
assert_eq!(history[0].state, WorkflowState::Completed);
}
#[test]
fn test_statistics() {
let service = MonitoringService::new();
let wf1 = WorkflowId::new();
let wf2 = WorkflowId::new();
service.start_workflow(wf1, "workflow1".to_string(), 2);
service.start_workflow(wf2, "workflow2".to_string(), 3);
service.complete_workflow(wf1, true);
service.complete_workflow(wf2, false);
let stats = service.get_statistics();
assert_eq!(stats.active_workflows, 0);
assert_eq!(stats.total_workflows, 2);
assert_eq!(stats.completed_workflows, 1);
assert_eq!(stats.failed_workflows, 1);
}
}