diff --git a/Cargo.toml b/Cargo.toml index 2bb69ad..16d4403 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ato" -version = "2.0.0" +version = "2.0.1" edition = "2021" license = "MIT" description = "A very minimal no-std async runtime" diff --git a/README.md b/README.md index 431ce24..d76a1e0 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ Add ATO to your `Cargo.toml`: ```toml [dependencies] -ato = "2.0.0" # Replace with the desired version +ato = "2.0.1" # Replace with the desired version ``` ## Usage @@ -47,13 +47,13 @@ const SPAWNER_SIZE: usize = 4; // Must be a power of two, e.g., 2, 4, 8, 16, etc fn main() { // create a spawner with the specified size - let spawner: Spawner = Spawner::new(); + let spawner: Spawner = Spawner::default(); // create a simple task that prints a message - let mut task = ato::task!({ - println!("Task 1 started"); + ato::task!(task, { + println!("Hello, World!"); }); - spawner.spawn(&mut task).unwrap(); + spawner.spawn(task).unwrap(); // run until all tasks are done running spawner.run_until_all_done().unwrap(); diff --git a/examples/simple.rs b/examples/simple.rs index 6f585cf..1a693a5 100644 --- a/examples/simple.rs +++ b/examples/simple.rs @@ -7,10 +7,10 @@ fn main() { let spawner: Spawner = Spawner::default(); // create a simple task that prints a message - let mut task = ato::task!({ - println!("Task 1 started"); + ato::task!(task, { + println!("Hello, World!"); }); - spawner.spawn(&mut task).unwrap(); + spawner.spawn(task).unwrap(); // run until all tasks are done running spawner.run_until_all_done().unwrap(); diff --git a/examples/time.rs b/examples/time.rs index d80ef6a..c75fc6b 100644 --- a/examples/time.rs +++ b/examples/time.rs @@ -12,7 +12,7 @@ fn get_platform_time() -> Duration { fn main() { let spawner: Spawner = Spawner::default(); - let mut task = ato::task!({ + ato::task!(task, { let start = Instant::now(); sleep(Duration::from_millis(200), get_platform_time).await; let elapsed = Instant::now().duration_since(start); @@ -21,7 +21,7 @@ fn main() { elapsed.as_millis() ); }); - spawner.spawn(&mut task).unwrap(); + spawner.spawn(task).unwrap(); spawner.run_until_all_done().unwrap(); } diff --git a/examples/until_condition.rs b/examples/until_condition.rs new file mode 100644 index 0000000..eba498e --- /dev/null +++ b/examples/until_condition.rs @@ -0,0 +1,32 @@ +use std::sync::atomic::AtomicBool; + +use ato::Spawner; + +const SPAWNER_SIZE: usize = 4; // Must be a power of two, e.g., 2, 4, 8, 16, etc. +static TASK_DONE: AtomicBool = AtomicBool::new(false); + +fn main() { + // create a spawner with the specified size + let spawner: Spawner = Spawner::default(); + + // create a simple task that prints a message + ato::task!(task, { + let mut i = 0; + while i < 5 { + println!("Running {}", i); + i += 1; + } + TASK_DONE.store(true, std::sync::atomic::Ordering::SeqCst); + }); + spawner.spawn(task).unwrap(); + + loop { + if TASK_DONE.load(std::sync::atomic::Ordering::SeqCst) { + break; + } + spawner.run_once().unwrap(); + } + + println!("Task completed."); + assert!(TASK_DONE.load(std::sync::atomic::Ordering::SeqCst)); +} diff --git a/examples/yield.rs b/examples/yield.rs index 3a2ce18..9396a07 100644 --- a/examples/yield.rs +++ b/examples/yield.rs @@ -8,7 +8,7 @@ fn main() { let spawner: Spawner = Spawner::default(); let lock = Arc::new(Mutex::new(Vec::new())); let lock_clone = lock.clone(); - let mut task = ato::task!({ + ato::task!(task, { { let mut num = lock_clone.lock().unwrap(); num.push(1); @@ -19,16 +19,16 @@ fn main() { num.push(3); } }); - spawner.spawn(&mut task).unwrap(); + spawner.spawn(task).unwrap(); let lock_clone = lock.clone(); - let mut task_2 = ato::task!({ + ato::task!(task_2, { { let mut num = lock_clone.lock().unwrap(); num.push(2); } }); - spawner.spawn(&mut task_2).unwrap(); + spawner.spawn(task_2).unwrap(); spawner.run_until_all_done().unwrap(); diff --git a/src/lib.rs b/src/lib.rs index 4c2837b..39d748a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -41,9 +41,21 @@ unsafe fn nop_clone(_data: *const ()) -> RawWaker { } static VTABLE: RawWakerVTable = RawWakerVTable::new(nop_clone, nop, nop, nop); -// Task type alias +/// A type alias for a pinned future that outputs `()`. type Task<'a> = Pin<&'a mut (dyn Future + Send + Sync)>; +/// A handle to a task (future) that can be spawned in the ATO runtime. +pub struct TaskHandle<'a> { + inner: Task<'a>, +} + +impl<'a> TaskHandle<'a> { + /// Creates a new TaskHandle from a future reference. + pub fn new(future: Pin<&'a mut (dyn Future + Send + Sync)>) -> Self { + TaskHandle { inner: future } + } +} + /// A simple task spawner and runner for `no_std` environments. /// The `Spawner` can queue and run multiple tasks (futures) in a FIFO manner. /// # Type Parameters @@ -66,15 +78,8 @@ impl<'a, const N: usize> Default for Spawner<'a, N> { impl<'a, const N: usize> Spawner<'a, N> { /// Spawns a task. Make sure to use the `task!` macro to pin the future to the stack. - pub fn spawn( - &self, - future: &'a mut (dyn Future + Send + Sync), - ) -> Result<(), Error> { - // We re-pin the reference. This is safe because the reference we receive - // is already mutable and valid for 'a. - let pinned_task = unsafe { Pin::new_unchecked(future) }; - - match self.tasks.enqueue(pinned_task) { + pub fn spawn(&self, future: TaskHandle<'a>) -> Result<(), Error> { + match self.tasks.enqueue(future.inner) { Ok(()) => Ok(()), Err(_) => Err(Error::TaskQueueFail), } @@ -96,6 +101,25 @@ impl<'a, const N: usize> Spawner<'a, N> { } Ok(()) } + + /// Runs a single task if available. + /// Could be useful in environments where you want to interleave task execution + /// with other processing. + pub fn run_once(&self) -> Result<(), Error> { + let mut cx = Context::from_waker(&self.waker); + + if let Some(mut task) = self.tasks.dequeue() { + match task.as_mut().poll(&mut cx) { + Poll::Ready(()) => {} + Poll::Pending => { + if self.tasks.enqueue(task).is_err() { + return Err(Error::TaskQueueFail); + } + } + } + } + Ok(()) + } } /// task! macro is used to create a pinned async block. @@ -103,8 +127,18 @@ impl<'a, const N: usize> Spawner<'a, N> { /// in the ATO runtime. #[macro_export] macro_rules! task { - ( $($body:tt)* ) => { - core::pin::pin!(async move { $($body)* }) + // Usage: task!(task_variable_name, { async_code... }); + ($name:ident, $body:expr) => { + // 1. Create the future variable in the CURRENT scope + let future = async move { $body }; + + // 2. Pin it. `core::pin::pin!` creates a `Pin<&mut T>` + // that borrows the local `future` variable we just created. + let pinned_fut = core::pin::pin!(future); + + // 3. Create the handle with the user-provided name. + // We assume TaskHandle::new is accessible here. + let $name = $crate::TaskHandle::new(pinned_fut); }; } @@ -152,12 +186,12 @@ mod tests { let sleep_duration = Duration::from_millis(10); - let mut pinned_future = task!({ + task!(pinned_future, { sleep(sleep_duration, get_current_test_time_duration).await; hello().await; }); - if let Err(_) = spawner.spawn(&mut pinned_future) { + if let Err(_) = spawner.spawn(pinned_future) { panic!("Failed to spawn task"); } @@ -172,7 +206,7 @@ mod tests { let spawner: Spawner<2> = Spawner::default(); let _ = get_test_epoch(); - let mut dequeue_future = task!({ + task!(dequeue_future, { loop { sleep(Duration::from_millis(10), get_current_test_time_duration).await; if let Some(_) = Q.dequeue() { @@ -181,17 +215,13 @@ mod tests { } }); - spawner - .spawn(&mut dequeue_future) - .expect("Failed to spawn task"); + spawner.spawn(dequeue_future).expect("Failed to spawn task"); - let mut enqueue_future = task!({ + task!(enqueue_future, { sleep(Duration::from_secs(1), get_current_test_time_duration).await; Q.enqueue(42).unwrap(); }); - spawner - .spawn(&mut enqueue_future) - .expect("Failed to spawn task"); + spawner.spawn(enqueue_future).expect("Failed to spawn task"); spawner.run_until_all_done().expect("Failed to run tasks"); } @@ -202,7 +232,7 @@ mod tests { let lock = Arc::new(Mutex::new(Vec::new())); let lock_clone = lock.clone(); - let mut first_future = task!({ + task!(first_future, { { let mut num = lock_clone.lock().unwrap(); num.push(1); @@ -215,7 +245,7 @@ mod tests { }); let lock_clone = lock.clone(); - let mut second_future = task!({ + task!(second_future, { { let mut num = lock_clone.lock().unwrap(); num.push(2); @@ -223,10 +253,10 @@ mod tests { }); spawner - .spawn(&mut first_future) + .spawn(first_future) .expect("Failed to spawn first future"); spawner - .spawn(&mut second_future) + .spawn(second_future) .expect("Failed to spawn second future"); spawner.run_until_all_done().unwrap();