Skip to main content

taskgraph/core/
task.rs

1//! Task definitions and storage traits.
2//! Handles both static function pointers and dynamic closures.
3
4use crate::error::Result;
5use core::time::Duration;
6
7#[cfg(feature = "alloc")]
8extern crate alloc;
9#[cfg(feature = "alloc")]
10use alloc::boxed::Box;
11
12#[cfg(feature = "alloc")]
13use core::pin::Pin;
14#[cfg(feature = "alloc")]
15use core::future::Future;
16
17/// Status of a task in the graph.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum TaskStatus {
20    Pending,
21    Running,
22    Completed,
23    Failed,
24    Retrying,
25}
26
27/// Type alias for async task logic to reduce complexity.
28#[cfg(all(feature = "alloc", feature = "async"))]
29pub type AsyncTask = Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result<()>> + Send>> + Send + Sync>;
30
31/// The body of a task.
32/// In `no_std` without `alloc`, only `fn()` pointers are supported.
33/// With `alloc`, `Box<dyn Fn() -> Result<()>>` is supported for closures.
34/// With `async`, `Box<dyn Fn() -> Pin<Box<dyn Future>>>` is supported for async tasks.
35pub enum TaskBody {
36    /// Static function pointer (no capture).
37    Static(fn() -> Result<()>),
38    /// Dynamic closure with capture.
39    #[cfg(feature = "alloc")]
40    Dynamic(Box<dyn Fn() -> Result<()> + Send + Sync>),
41    /// Async task (requires alloc and async).
42    #[cfg(all(feature = "alloc", feature = "async"))]
43    Async(AsyncTask),
44}
45
46/// A task in the DAG.
47pub struct Task {
48    /// Human-readable name of the task.
49    pub name: &'static str,
50    /// The executable logic.
51    pub body: TaskBody,
52    /// Number of retries allowed.
53    pub retries: u32,
54    /// Timeout duration.
55    pub timeout: Option<Duration>,
56    /// Current execution status.
57    pub status: TaskStatus,
58    /// Number of remaining dependencies before this task can run.
59    pub remaining_deps: core::sync::atomic::AtomicUsize,
60    /// Initial number of dependencies.
61    pub initial_deps: usize,
62    /// Current number of retries attempted.
63    pub current_retry: core::sync::atomic::AtomicU32,
64}
65
66impl Task {
67    /// Create a new static task.
68    pub fn new_static(name: &'static str, body: fn() -> Result<()>) -> Self {
69        Self {
70            name,
71            body: TaskBody::Static(body),
72            retries: 0,
73            timeout: None,
74            status: TaskStatus::Pending,
75            remaining_deps: core::sync::atomic::AtomicUsize::new(0),
76            initial_deps: 0,
77            current_retry: core::sync::atomic::AtomicU32::new(0),
78        }
79    }
80
81    /// Create a new dynamic task (requires `alloc`).
82    #[cfg(feature = "alloc")]
83    pub fn new_dynamic<F>(name: &'static str, body: F) -> Self 
84    where F: Fn() -> Result<()> + Send + Sync + 'static
85    {
86        Self {
87            name,
88            body: TaskBody::Dynamic(Box::new(body)),
89            retries: 0,
90            timeout: None,
91            status: TaskStatus::Pending,
92            remaining_deps: core::sync::atomic::AtomicUsize::new(0),
93            initial_deps: 0,
94            current_retry: core::sync::atomic::AtomicU32::new(0),
95        }
96    }
97
98    /// Create a new async task (requires alloc and async).
99    #[cfg(all(feature = "alloc", feature = "async"))]
100    pub fn new_async<F>(name: &'static str, body: F) -> Self 
101    where F: Fn() -> Pin<Box<dyn Future<Output = Result<()>> + Send>> + Send + Sync + 'static
102    {
103        Self {
104            name,
105            body: TaskBody::Async(Box::new(body)),
106            retries: 0,
107            timeout: None,
108            status: TaskStatus::Pending,
109            remaining_deps: core::sync::atomic::AtomicUsize::new(0),
110            initial_deps: 0,
111            current_retry: core::sync::atomic::AtomicU32::new(0),
112        }
113    }
114
115    /// Set the number of retries for this task.
116    pub fn with_retries(mut self, retries: u32) -> Self {
117        self.retries = retries;
118        self
119    }
120
121    /// Set a timeout for this task.
122    pub fn with_timeout(mut self, timeout: Duration) -> Self {
123        self.timeout = Some(timeout);
124        self
125    }
126}
127
128/// Interface for task storage backends.
129pub trait TaskStore {
130    /// Add a task to the storage.
131    fn add_task(&mut self, task: Task, deps: &[usize]) -> Result<usize>;
132    /// Get tasks that depend on this one (outgoing edges).
133    fn get_successors(&self, id: usize) -> u64;
134    /// Get a task by its ID.
135    fn get_task(&self, id: usize) -> Option<&Task>;
136    /// Mutably get a task by its ID.
137    fn get_task_mut(&mut self, id: usize) -> Option<&mut Task>;
138    /// Update the status of a task.
139    fn update_status(&mut self, id: usize, status: TaskStatus) -> Result<()>;
140    /// Get the total number of tasks.
141    fn task_count(&self) -> usize;
142}