130 lines
3.6 KiB
Rust
130 lines
3.6 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Mutex;
|
|
|
|
use chrono::{NaiveDateTime, Utc};
|
|
use serde::Serialize;
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct TaskProgress {
|
|
pub current: u64,
|
|
pub total: u64,
|
|
pub message: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct TaskInfo {
|
|
pub id: String,
|
|
pub task_type: String,
|
|
pub status: TaskStatus,
|
|
pub progress: Option<TaskProgress>,
|
|
pub started_at: NaiveDateTime,
|
|
pub completed_at: Option<NaiveDateTime>,
|
|
pub result: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
|
pub enum TaskStatus {
|
|
Pending,
|
|
Running,
|
|
Completed,
|
|
Failed,
|
|
}
|
|
|
|
pub struct TaskManager {
|
|
tasks: Mutex<HashMap<String, TaskInfo>>,
|
|
}
|
|
|
|
impl Default for TaskManager {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl TaskManager {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
tasks: Mutex::new(HashMap::new()),
|
|
}
|
|
}
|
|
|
|
/// Register a new task as Running. Returns the task ID.
|
|
pub fn register(&self, task_type: &str) -> String {
|
|
let id = uuid::Uuid::new_v4().to_string();
|
|
let info = TaskInfo {
|
|
id: id.clone(),
|
|
task_type: task_type.to_string(),
|
|
status: TaskStatus::Running,
|
|
progress: None,
|
|
started_at: Utc::now().naive_utc(),
|
|
completed_at: None,
|
|
result: None,
|
|
};
|
|
self.tasks.lock().unwrap().insert(id.clone(), info);
|
|
id
|
|
}
|
|
|
|
/// Register a new task as Pending (queued, not yet running). Returns the task ID.
|
|
pub fn register_pending(&self, task_type: &str) -> String {
|
|
let id = uuid::Uuid::new_v4().to_string();
|
|
let info = TaskInfo {
|
|
id: id.clone(),
|
|
task_type: task_type.to_string(),
|
|
status: TaskStatus::Pending,
|
|
progress: None,
|
|
started_at: Utc::now().naive_utc(),
|
|
completed_at: None,
|
|
result: None,
|
|
};
|
|
self.tasks.lock().unwrap().insert(id.clone(), info);
|
|
id
|
|
}
|
|
|
|
/// Transition a task from Pending to Running.
|
|
pub fn start(&self, id: &str) {
|
|
if let Some(task) = self.tasks.lock().unwrap().get_mut(id) {
|
|
task.status = TaskStatus::Running;
|
|
}
|
|
}
|
|
|
|
/// Update progress on a running task.
|
|
pub fn update_progress(&self, id: &str, current: u64, total: u64, message: &str) {
|
|
if let Some(task) = self.tasks.lock().unwrap().get_mut(id) {
|
|
task.progress = Some(TaskProgress {
|
|
current,
|
|
total,
|
|
message: message.to_string(),
|
|
});
|
|
}
|
|
}
|
|
|
|
/// Mark a task as completed with a result string.
|
|
pub fn complete(&self, id: &str, result: String) {
|
|
if let Some(task) = self.tasks.lock().unwrap().get_mut(id) {
|
|
task.status = TaskStatus::Completed;
|
|
task.progress = None;
|
|
task.completed_at = Some(Utc::now().naive_utc());
|
|
task.result = Some(result);
|
|
}
|
|
}
|
|
|
|
/// Mark a task as failed with an error message.
|
|
pub fn fail(&self, id: &str, error: String) {
|
|
if let Some(task) = self.tasks.lock().unwrap().get_mut(id) {
|
|
task.status = TaskStatus::Failed;
|
|
task.progress = None;
|
|
task.completed_at = Some(Utc::now().naive_utc());
|
|
task.result = Some(error);
|
|
}
|
|
}
|
|
|
|
/// Get a task by ID.
|
|
pub fn get(&self, id: &str) -> Option<TaskInfo> {
|
|
self.tasks.lock().unwrap().get(id).cloned()
|
|
}
|
|
|
|
/// List all tasks.
|
|
pub fn list(&self) -> Vec<TaskInfo> {
|
|
self.tasks.lock().unwrap().values().cloned().collect()
|
|
}
|
|
}
|