use crate::error::{Result, TaskFlowError};
use crate::executor::{
Executor,
handlers::{
FileTaskHandler, HttpTaskHandler, NotificationTaskHandler, PythonTaskHandler,
ShellTaskHandler,
},
};
use crate::scheduler::Scheduler;
use crate::storage::{InMemoryStorage, TaskStorage};
use crate::task::{TaskDefinition, TaskHandler, TaskStatus};
use std::path::Path;
use std::sync::Arc;
use tracing::{error, info};
pub mod config;
pub mod metrics;
pub mod yaml_config;
pub use config::{StorageType, TaskFlowConfig};
pub use metrics::TaskMetrics;
pub use yaml_config::{load_from_yaml_file, load_from_yaml_str};
pub struct TaskFlow {
scheduler: Arc<Scheduler>,
executor: Arc<tokio::sync::Mutex<Executor>>,
storage: Arc<dyn TaskStorage>,
}
impl TaskFlow {
pub async fn new(config: TaskFlowConfig) -> Result<Self> {
let storage: Arc<dyn TaskStorage> = match config.storage_type {
StorageType::InMemory => {
info!("Using in-memory storage");
Arc::new(InMemoryStorage::new())
}
};
let scheduler = Arc::new(Scheduler::new(Arc::clone(&storage), config.scheduler));
let mut executor = Executor::new(Arc::clone(&scheduler), config.executor);
executor.register_handler(Arc::new(HttpTaskHandler::new()));
executor.register_handler(Arc::new(ShellTaskHandler::new()));
executor.register_handler(Arc::new(PythonTaskHandler::new()));
executor.register_handler(Arc::new(FileTaskHandler::new()));
executor.register_handler(Arc::new(NotificationTaskHandler::new()));
Ok(Self {
scheduler,
executor: Arc::new(tokio::sync::Mutex::new(executor)),
storage,
})
}
pub async fn from_yaml_file<P: AsRef<Path>>(path: P) -> Result<Self> {
let config = load_from_yaml_file(path)?;
Self::new(config).await
}
pub async fn from_yaml_str(config_content: &str) -> Result<Self> {
let config = load_from_yaml_str(config_content)?;
Self::new(config).await
}
pub async fn register_handler(&self, handler: Arc<dyn TaskHandler>) {
let mut executor = self.executor.lock().await;
executor.register_handler(handler);
}
pub async fn submit_task(&self, definition: TaskDefinition) -> Result<String> {
self.scheduler.submit_task(definition).await
}
pub async fn submit_http_task(
&self,
name: &str,
url: &str,
method: Option<&str>,
) -> Result<String> {
let mut definition = TaskDefinition::new(name, "http_request")
.with_payload("url", serde_json::Value::String(url.to_string()));
if let Some(method) = method {
definition =
definition.with_payload("method", serde_json::Value::String(method.to_string()));
}
self.submit_task(definition).await
}
pub async fn submit_shell_task(
&self,
name: &str,
command: &str,
args: Vec<&str>,
) -> Result<String> {
let args_json: Vec<serde_json::Value> = args
.into_iter()
.map(|arg| serde_json::Value::String(arg.to_string()))
.collect();
let definition = TaskDefinition::new(name, "shell_command")
.with_payload("command", serde_json::Value::String(command.to_string()))
.with_payload("args", serde_json::Value::Array(args_json));
self.submit_task(definition).await
}
pub async fn get_task_status(&self, task_id: &str) -> Result<Option<TaskStatus>> {
self.scheduler.get_task_status(task_id).await
}
pub async fn cancel_task(&self, task_id: &str) -> Result<()> {
self.scheduler.cancel_task(task_id).await
}
pub async fn list_tasks(&self, status: Option<TaskStatus>) -> Result<Vec<crate::task::Task>> {
self.scheduler.list_tasks(status).await
}
pub async fn start(&self) -> Result<()> {
info!("Starting TaskFlow framework");
let scheduler_clone = Arc::clone(&self.scheduler);
let scheduler_task = tokio::spawn(async move {
if let Err(e) = scheduler_clone.start().await {
error!("Scheduler failed: {}", e);
}
});
let executor_clone = Arc::clone(&self.executor);
let executor_task = tokio::spawn(async move {
let executor = executor_clone.lock().await;
if let Err(e) = executor.start().await {
error!("Executor failed: {}", e);
}
});
tokio::select! {
result = scheduler_task => {
if let Err(e) = result {
error!("Scheduler task panicked: {}", e);
}
}
result = executor_task => {
if let Err(e) = result {
error!("Executor task panicked: {}", e);
}
}
}
Ok(())
}
pub async fn shutdown(&self) -> Result<()> {
info!("Shutting down TaskFlow framework");
let executor = self.executor.lock().await;
executor.shutdown().await;
Ok(())
}
pub async fn wait_for_completion(
&self,
task_id: &str,
timeout_seconds: Option<u64>,
) -> Result<crate::task::Task> {
let timeout_duration = timeout_seconds.unwrap_or(300);
let start_time = std::time::Instant::now();
loop {
if let Some(task) = self.storage.get_task(task_id).await? {
if task.is_finished() {
return Ok(task);
}
} else {
return Err(TaskFlowError::TaskNotFound(task_id.to_string()));
}
if start_time.elapsed().as_secs() > timeout_duration {
return Err(TaskFlowError::TimeoutError);
}
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
}
}
pub async fn get_task_metrics(&self) -> Result<TaskMetrics> {
let pending = self.list_tasks(Some(TaskStatus::Pending)).await?.len();
let running = self.list_tasks(Some(TaskStatus::Running)).await?.len();
let completed = self.list_tasks(Some(TaskStatus::Completed)).await?.len();
let failed = self.list_tasks(Some(TaskStatus::Failed)).await?.len();
let cancelled = self.list_tasks(Some(TaskStatus::Cancelled)).await?.len();
Ok(TaskMetrics {
pending,
running,
completed,
failed,
cancelled,
total: pending + running + completed + failed + cancelled,
})
}
}