use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Condvar, Mutex}; use std::time::{Duration, Instant}; /// Unique task identifier. pub type TaskId = u64; /// Task status. #[derive(Debug, Clone, PartialEq, Eq)] pub enum TaskStatus { Pending, Running, Completed, Failed(String), Cancelled, } /// Immutable task state exposed to UI consumers. #[derive(Debug, Clone)] pub struct TaskSnapshot { pub id: TaskId, pub label: String, pub group_id: Option, pub group_name: Option, pub status: TaskStatus, pub progress: Option, pub message: Option, pub created_at: Instant, } /// Entry tracking a task. #[derive(Debug)] struct TaskEntry { id: TaskId, label: String, group_id: Option, group_name: Option, status: TaskStatus, cancel_flag: Arc, progress: Option, message: Option, created_at: Instant, finished_at: Option, last_progress_report: Option, } /// Manages concurrent tasks with a max concurrency limit and FIFO queue. pub struct TaskManager { max_concurrent: usize, next_id: Mutex, tasks: Mutex>, state_changed: Condvar, } impl TaskManager { /// Create a new task manager with the given concurrency limit. pub fn new(max_concurrent: usize) -> Self { Self { max_concurrent, next_id: Mutex::new(1), tasks: Mutex::new(Vec::new()), state_changed: Condvar::new(), } } /// Submit a new task. Returns its unique identifier. pub fn submit(&self, label: &str) -> TaskId { let mut next = self.next_id.lock().unwrap(); let id = *next; *next += 1; let entry = TaskEntry { id, label: label.to_owned(), group_id: None, group_name: None, status: TaskStatus::Pending, cancel_flag: Arc::new(AtomicBool::new(false)), progress: None, message: None, created_at: Instant::now(), finished_at: None, last_progress_report: None, }; let mut tasks = self.tasks.lock().unwrap(); tasks.push(entry); // Auto-start if under capacity let running = tasks .iter() .filter(|t| t.status == TaskStatus::Running) .count(); if running < self.max_concurrent && let Some(t) = tasks .iter_mut() .find(|t| t.id == id && t.status == TaskStatus::Pending) { t.status = TaskStatus::Running; } id } /// Submit a new task within a group. Returns its unique identifier. pub fn submit_grouped(&self, label: &str, group_id: &str, group_name: &str) -> TaskId { let id = self.submit(label); let mut tasks = self.tasks.lock().unwrap(); if let Some(entry) = tasks.iter_mut().find(|t| t.id == id) { entry.group_id = Some(group_id.to_owned()); entry.group_name = Some(group_name.to_owned()); } id } /// Block a worker until its task may run. Returns false if cancelled. pub fn wait_until_runnable(&self, task_id: TaskId) -> bool { let mut tasks = self.tasks.lock().unwrap(); loop { match tasks .iter() .find(|task| task.id == task_id) .map(|task| &task.status) { Some(TaskStatus::Running) => return true, Some(TaskStatus::Pending) => tasks = self.state_changed.wait(tasks).unwrap(), _ => return false, } } } /// Mark a task as completed. pub fn complete(&self, task_id: TaskId) { let mut tasks = self.tasks.lock().unwrap(); if let Some(entry) = tasks.iter_mut().find(|t| t.id == task_id) && matches!(entry.status, TaskStatus::Running) { entry.status = TaskStatus::Completed; entry.progress = Some(1.0); entry.finished_at = Some(Instant::now()); } Self::promote_next(&mut tasks, self.max_concurrent); self.state_changed.notify_all(); } /// Mark a task as failed with an error message. pub fn fail(&self, task_id: TaskId, error: String) { let mut tasks = self.tasks.lock().unwrap(); if let Some(entry) = tasks.iter_mut().find(|t| t.id == task_id) && matches!(entry.status, TaskStatus::Running) { entry.message = Some(error.clone()); entry.status = TaskStatus::Failed(error); entry.finished_at = Some(Instant::now()); } Self::promote_next(&mut tasks, self.max_concurrent); self.state_changed.notify_all(); } /// Cancel a task by setting its cancel flag and status. pub fn cancel(&self, task_id: TaskId) { let mut tasks = self.tasks.lock().unwrap(); if let Some(entry) = tasks.iter_mut().find(|t| t.id == task_id) && matches!(entry.status, TaskStatus::Running | TaskStatus::Pending) { entry.cancel_flag.store(true, Ordering::Release); entry.status = TaskStatus::Cancelled; entry.finished_at = Some(Instant::now()); } Self::promote_next(&mut tasks, self.max_concurrent); self.state_changed.notify_all(); } /// Cancel every active task in a group and release their workers. pub fn cancel_group(&self, group_id: &str) { let mut tasks = self.tasks.lock().unwrap(); let now = Instant::now(); for entry in tasks.iter_mut().filter(|task| { task.group_id.as_deref() == Some(group_id) && matches!(task.status, TaskStatus::Running | TaskStatus::Pending) }) { entry.cancel_flag.store(true, Ordering::Release); entry.status = TaskStatus::Cancelled; entry.finished_at = Some(now); } Self::promote_next(&mut tasks, self.max_concurrent); self.state_changed.notify_all(); } /// Return the group containing a task, if any. pub fn group_id(&self, task_id: TaskId) -> Option { self.tasks .lock() .unwrap() .iter() .find(|task| task.id == task_id) .and_then(|task| task.group_id.clone()) } /// Check whether a task has been cancelled. pub fn is_cancelled(&self, task_id: TaskId) -> bool { let tasks = self.tasks.lock().unwrap(); tasks .iter() .find(|t| t.id == task_id) .map(|t| t.cancel_flag.load(Ordering::Acquire)) .unwrap_or(false) } /// Shared cancellation flag for a worker owned by this task. pub fn cancellation_flag(&self, task_id: TaskId) -> Option> { self.tasks .lock() .unwrap() .iter() .find(|task| task.id == task_id) .map(|task| Arc::clone(&task.cancel_flag)) } /// Return the current status of a task. pub fn status(&self, task_id: TaskId) -> Option { let tasks = self.tasks.lock().unwrap(); tasks .iter() .find(|t| t.id == task_id) .map(|t| t.status.clone()) } /// Count tasks that are still queued. pub fn pending_count(&self) -> usize { let tasks = self.tasks.lock().unwrap(); tasks .iter() .filter(|t| t.status == TaskStatus::Pending) .count() } /// Count tasks that are currently running. pub fn running_count(&self) -> usize { let tasks = self.tasks.lock().unwrap(); tasks .iter() .filter(|t| t.status == TaskStatus::Running) .count() } /// Remove finished tasks older than the configured retention period. pub fn evict_expired(&self) { let cutoff = Instant::now() - FINISHED_TASK_TTL; let mut tasks = self.tasks.lock().unwrap(); tasks.retain(|task| { task.finished_at .is_none_or(|finished_at| finished_at > cutoff) }); } /// Remove every finished task while preserving running and queued work. pub fn clear_completed(&self) { self.tasks .lock() .unwrap() .retain(|task| matches!(task.status, TaskStatus::Pending | TaskStatus::Running)); } /// Update progress for a running task. Throttled to at most once per 250ms. pub fn report_progress(&self, task_id: TaskId, progress: Option, message: Option) { let mut tasks = self.tasks.lock().unwrap(); if let Some(entry) = tasks.iter_mut().find(|t| t.id == task_id) && entry.status == TaskStatus::Running { let now = Instant::now(); let should_report = match entry.last_progress_report { Some(prev) => now.duration_since(prev).as_millis() >= PROGRESS_THROTTLE_MS as u128, None => true, }; if should_report { entry.progress = progress; entry.message = message; entry.last_progress_report = Some(now); } } } /// Return the current progress of a task. pub fn progress(&self, task_id: TaskId) -> Option { let tasks = self.tasks.lock().unwrap(); tasks .iter() .find(|t| t.id == task_id) .and_then(|t| t.progress) } /// Return a snapshot of all tasks for UI display. pub fn snapshots(&self) -> Vec { let tasks = self.tasks.lock().unwrap(); let mut snapshots = tasks .iter() .filter(|task| task.finished_at.is_none()) .chain( tasks .iter() .rev() .filter(|task| task.finished_at.is_some()) .take(RECENT_FINISHED_LIMIT), ) .map(|task| TaskSnapshot { id: task.id, label: task.label.clone(), group_id: task.group_id.clone(), group_name: task.group_name.clone(), status: task.status.clone(), progress: task.progress, message: task.message.clone(), created_at: task.created_at, }) .collect::>(); snapshots.sort_by_key(|snapshot| snapshot.created_at); snapshots } /// Promote the next queued task to running if capacity allows. fn promote_next(tasks: &mut [TaskEntry], max_concurrent: usize) { while tasks .iter() .filter(|task| task.status == TaskStatus::Running) .count() < max_concurrent { let Some(task) = tasks .iter_mut() .find(|task| task.status == TaskStatus::Pending) else { break; }; task.status = TaskStatus::Running; } } } impl Default for TaskManager { fn default() -> Self { Self::new(3) } } /// Default progress throttle interval (250ms per spec). pub const PROGRESS_THROTTLE_MS: u64 = 250; pub const RECENT_FINISHED_LIMIT: usize = 10; pub const FINISHED_TASK_TTL: Duration = Duration::from_secs(60 * 60); #[cfg(test)] mod tests { use super::*; #[test] fn submit_and_start() { let mgr = TaskManager::default(); let id = mgr.submit("build site"); // Auto-started since capacity allows assert_eq!(mgr.status(id), Some(TaskStatus::Running)); } #[test] fn max_concurrent_enforced() { let mgr = TaskManager::new(3); let ids: Vec = (0..4).map(|i| mgr.submit(&format!("task {i}"))).collect(); // First 3 auto-started, 4th stays queued assert_eq!(mgr.running_count(), 3); assert_eq!(mgr.status(ids[3]), Some(TaskStatus::Pending)); } #[test] fn fifo_order() { let mgr = TaskManager::new(1); // limit to 1 to test FIFO let a = mgr.submit("first"); // auto-starts let b = mgr.submit("second"); // queued let c = mgr.submit("third"); // queued assert_eq!(mgr.status(a), Some(TaskStatus::Running)); mgr.complete(a); // should auto-promote b assert_eq!(mgr.status(b), Some(TaskStatus::Running)); assert_eq!(mgr.status(c), Some(TaskStatus::Pending)); } #[test] fn cancel_sets_flag() { let mgr = TaskManager::default(); let id = mgr.submit("upload"); // Task is auto-started (Running) assert!(!mgr.is_cancelled(id)); mgr.cancel(id); assert!(mgr.is_cancelled(id)); assert_eq!(mgr.status(id), Some(TaskStatus::Cancelled)); } #[test] fn complete_and_fail() { let mgr = TaskManager::default(); let ok = mgr.submit("good task"); let bad = mgr.submit("bad task"); // Both auto-started (capacity=3) mgr.complete(ok); mgr.fail(bad, "disk full".into()); assert_eq!(mgr.status(ok), Some(TaskStatus::Completed)); assert_eq!( mgr.status(bad), Some(TaskStatus::Failed("disk full".into())) ); // Progress should be 1.0 on completed assert_eq!(mgr.progress(ok), Some(1.0)); } #[test] fn eviction_removes_only_expired_finished_tasks() { let mgr = TaskManager::new(3); let a = mgr.submit("done"); // auto-starts let b = mgr.submit("broken"); // auto-starts let e = mgr.submit("busy"); // auto-starts let _c = mgr.submit("stopped"); // queued (at capacity) let _d = mgr.submit("waiting"); // queued mgr.complete(a); mgr.fail(b, "oops".into()); // c should have been auto-promoted when a completed, and again when b failed // After a completes: c promoted to running // After b fails: d promoted to running { let mut tasks = mgr.tasks.lock().unwrap(); for task in tasks.iter_mut().filter(|task| task.finished_at.is_some()) { task.finished_at = Some(Instant::now() - FINISHED_TASK_TTL - Duration::from_secs(1)); } } mgr.evict_expired(); assert_eq!(mgr.status(a), None); assert_eq!(mgr.status(b), None); // c and d were promoted, e is still running assert_eq!(mgr.status(e), Some(TaskStatus::Running)); } #[test] fn snapshots_retain_only_ten_finished_tasks() { let mgr = TaskManager::new(20); for index in 0..12 { let id = mgr.submit(&format!("task {index}")); mgr.complete(id); } assert_eq!(mgr.snapshots().len(), 10); } #[test] fn clear_completed_preserves_active_tasks() { let mgr = TaskManager::new(2); let done = mgr.submit("done"); let running = mgr.submit("running"); mgr.complete(done); mgr.clear_completed(); assert_eq!(mgr.status(done), None); assert_eq!(mgr.status(running), Some(TaskStatus::Running)); } #[test] fn completing_task_starts_next_queued() { let mgr = TaskManager::new(1); let a = mgr.submit("first"); // auto-starts let b = mgr.submit("second"); // queued assert_eq!(mgr.status(a), Some(TaskStatus::Running)); assert_eq!(mgr.status(b), Some(TaskStatus::Pending)); mgr.complete(a); assert_eq!(mgr.status(b), Some(TaskStatus::Running)); } #[test] fn queued_task_waits_for_a_slot() { let mgr = std::sync::Arc::new(TaskManager::new(1)); let running = mgr.submit("running"); let queued = mgr.submit("queued"); let waiter = { let mgr = mgr.clone(); std::thread::spawn(move || mgr.wait_until_runnable(queued)) }; assert!(!waiter.is_finished()); mgr.complete(running); assert!(waiter.join().unwrap()); } #[test] fn cancelling_queued_task_stops_its_waiter() { let mgr = std::sync::Arc::new(TaskManager::new(1)); let _running = mgr.submit("running"); let queued = mgr.submit("queued"); let waiter = { let mgr = mgr.clone(); std::thread::spawn(move || mgr.wait_until_runnable(queued)) }; mgr.cancel(queued); assert!(!waiter.join().unwrap()); } #[test] fn cancel_precondition_ignores_completed() { let mgr = TaskManager::default(); let id = mgr.submit("task"); mgr.complete(id); mgr.cancel(id); // should be no-op assert_eq!(mgr.status(id), Some(TaskStatus::Completed)); } #[test] fn cancelling_group_settles_every_active_task() { let mgr = TaskManager::new(2); let first = mgr.submit_grouped("first", "generation-1", "Render Site"); let second = mgr.submit_grouped("second", "generation-1", "Render Site"); let third = mgr.submit_grouped("third", "generation-1", "Render Site"); let unrelated = mgr.submit("unrelated"); mgr.cancel_group("generation-1"); assert_eq!(mgr.status(first), Some(TaskStatus::Cancelled)); assert_eq!(mgr.status(second), Some(TaskStatus::Cancelled)); assert_eq!(mgr.status(third), Some(TaskStatus::Cancelled)); assert_ne!(mgr.status(unrelated), Some(TaskStatus::Cancelled)); } #[test] fn report_progress_updates_task() { let mgr = TaskManager::default(); let id = mgr.submit("upload"); mgr.report_progress(id, Some(0.5), Some("halfway".into())); assert_eq!(mgr.progress(id), Some(0.5)); assert_eq!(mgr.snapshots()[0].message.as_deref(), Some("halfway")); } }