use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use taskflow_rs::{
Task, TaskDefinition, TaskFlow, TaskResult, error::TaskFlowError, framework::TaskFlowConfig,
task::TaskHandler,
};
use tracing_subscriber::fmt::init;
struct MathTaskHandler;
#[async_trait]
impl TaskHandler for MathTaskHandler {
async fn execute(&self, task: &Task) -> Result<TaskResult, TaskFlowError> {
let start_time = std::time::Instant::now();
let operation = task
.definition
.payload
.get("operation")
.and_then(|v| v.as_str())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing 'operation' in payload".to_string())
})?;
let a = task
.definition
.payload
.get("a")
.and_then(|v| v.as_f64())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing 'a' in payload".to_string())
})?;
let b = task
.definition
.payload
.get("b")
.and_then(|v| v.as_f64())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing 'b' in payload".to_string())
})?;
let result = match operation {
"add" => a + b,
"subtract" => a - b,
"multiply" => a * b,
"divide" => {
if b == 0.0 {
return Ok(TaskResult {
success: false,
output: None,
error: Some("Division by zero".to_string()),
execution_time_ms: start_time.elapsed().as_millis() as u64,
metadata: HashMap::new(),
});
}
a / b
}
_ => {
return Ok(TaskResult {
success: false,
output: None,
error: Some(format!("Unknown operation: {}", operation)),
execution_time_ms: start_time.elapsed().as_millis() as u64,
metadata: HashMap::new(),
});
}
};
let execution_time = start_time.elapsed().as_millis() as u64;
let mut metadata = HashMap::new();
metadata.insert("operation".to_string(), operation.to_string());
metadata.insert("operand_a".to_string(), a.to_string());
metadata.insert("operand_b".to_string(), b.to_string());
Ok(TaskResult {
success: true,
output: Some(result.to_string()),
error: None,
execution_time_ms: execution_time,
metadata,
})
}
fn task_type(&self) -> &str {
"math_operation"
}
}
struct DataProcessingHandler;
#[async_trait]
impl TaskHandler for DataProcessingHandler {
async fn execute(&self, task: &Task) -> Result<TaskResult, TaskFlowError> {
let start_time = std::time::Instant::now();
let data = task
.definition
.payload
.get("data")
.and_then(|v| v.as_array())
.ok_or_else(|| {
TaskFlowError::InvalidConfiguration("Missing 'data' array in payload".to_string())
})?;
let operation = task
.definition
.payload
.get("operation")
.and_then(|v| v.as_str())
.unwrap_or("sum");
let numbers: Result<Vec<f64>, _> = data
.iter()
.map(|v| {
v.as_f64().ok_or_else(|| {
TaskFlowError::InvalidConfiguration(
"All data elements must be numbers".to_string(),
)
})
})
.collect();
let numbers = numbers?;
let result = match operation {
"sum" => numbers.iter().sum::<f64>(),
"average" => {
if numbers.is_empty() {
0.0
} else {
numbers.iter().sum::<f64>() / numbers.len() as f64
}
}
"max" => numbers.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)),
"min" => numbers.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
_ => {
return Ok(TaskResult {
success: false,
output: None,
error: Some(format!("Unknown operation: {}", operation)),
execution_time_ms: start_time.elapsed().as_millis() as u64,
metadata: HashMap::new(),
});
}
};
let execution_time = start_time.elapsed().as_millis() as u64;
let mut metadata = HashMap::new();
metadata.insert("operation".to_string(), operation.to_string());
metadata.insert("data_count".to_string(), numbers.len().to_string());
Ok(TaskResult {
success: true,
output: Some(result.to_string()),
error: None,
execution_time_ms: execution_time,
metadata,
})
}
fn task_type(&self) -> &str {
"data_processing"
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init();
let config = TaskFlowConfig::with_in_memory();
let taskflow = TaskFlow::new(config).await?;
taskflow.register_handler(Arc::new(MathTaskHandler)).await;
taskflow
.register_handler(Arc::new(DataProcessingHandler))
.await;
println!("TaskFlow with custom handlers started!");
let add_task = TaskDefinition::new("addition", "math_operation")
.with_payload("operation", serde_json::Value::String("add".to_string()))
.with_payload("a", serde_json::Value::Number(serde_json::Number::from(10)))
.with_payload("b", serde_json::Value::Number(serde_json::Number::from(5)));
let add_task_id = taskflow.submit_task(add_task).await?;
println!("Submitted addition task: {}", add_task_id);
let multiply_task = TaskDefinition::new("multiplication", "math_operation")
.with_payload(
"operation",
serde_json::Value::String("multiply".to_string()),
)
.with_payload("a", serde_json::Value::Number(serde_json::Number::from(7)))
.with_payload("b", serde_json::Value::Number(serde_json::Number::from(3)));
let multiply_task_id = taskflow.submit_task(multiply_task).await?;
println!("Submitted multiplication task: {}", multiply_task_id);
let data_array = vec![
serde_json::Value::Number(serde_json::Number::from(1)),
serde_json::Value::Number(serde_json::Number::from(2)),
serde_json::Value::Number(serde_json::Number::from(3)),
serde_json::Value::Number(serde_json::Number::from(4)),
serde_json::Value::Number(serde_json::Number::from(5)),
];
let sum_task = TaskDefinition::new("sum_data", "data_processing")
.with_payload("operation", serde_json::Value::String("sum".to_string()))
.with_payload("data", serde_json::Value::Array(data_array.clone()));
let sum_task_id = taskflow.submit_task(sum_task).await?;
println!("Submitted sum task: {}", sum_task_id);
let avg_task = TaskDefinition::new("average_data", "data_processing")
.with_payload(
"operation",
serde_json::Value::String("average".to_string()),
)
.with_payload("data", serde_json::Value::Array(data_array))
.with_dependencies(vec![sum_task_id.clone()]);
let avg_task_id = taskflow.submit_task(avg_task).await?;
println!("Submitted average task (depends on sum): {}", avg_task_id);
let taskflow_clone = std::sync::Arc::new(taskflow);
let taskflow_for_execution = taskflow_clone.clone();
let execution_handle = tokio::spawn(async move {
if let Err(e) = taskflow_for_execution.start().await {
eprintln!("TaskFlow execution failed: {}", e);
}
});
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
loop {
let metrics = taskflow_clone.get_task_metrics().await?;
println!(
"Task metrics: pending={}, running={}, completed={}, failed={}",
metrics.pending, metrics.running, metrics.completed, metrics.failed
);
if metrics.pending == 0 && metrics.running == 0 {
break;
}
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
}
println!("\nAll tasks completed! Results:");
let tasks = taskflow_clone.list_tasks(None).await?;
for task in tasks {
println!(
"\nTask: {} ({})",
task.definition.name, task.definition.task_type
);
println!(" Status: {:?}", task.status);
if let Some(result) = &task.result {
if result.success {
println!(
" Result: {}",
result.output.as_ref().unwrap_or(&"No output".to_string())
);
println!(" Execution time: {}ms", result.execution_time_ms);
if !result.metadata.is_empty() {
println!(" Metadata: {:?}", result.metadata);
}
} else {
println!(
" Error: {}",
result
.error
.as_ref()
.unwrap_or(&"Unknown error".to_string())
);
}
}
}
let final_metrics = taskflow_clone.get_task_metrics().await?;
println!("\nFinal metrics:");
println!(" Total tasks: {}", final_metrics.total);
println!(
" Success rate: {:.1}%",
final_metrics.success_rate() * 100.0
);
taskflow_clone.shutdown().await?;
execution_handle.abort();
Ok(())
}