1 // Copyright 2024, The Android Open Source Project 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 //! This library implements a simple executor using cyclic scheduling. 16 17 #![cfg_attr(not(test), no_std)] 18 19 extern crate alloc; 20 use alloc::{boxed::Box, vec::Vec}; 21 use core::{future::Future, pin::Pin}; 22 use gbl_async::poll; 23 24 /// `CyclicExecutor` is a simple single thread executor that simply cyclically polls all Futures. 25 #[derive(Default)] 26 pub struct CyclicExecutor<'a> { 27 tasks: Vec<Pin<Box<dyn Future<Output = ()> + 'a>>>, 28 } 29 30 impl<'a> CyclicExecutor<'a> { 31 /// Adds a new task. spawn_task(&mut self, task: impl Future<Output = ()> + 'a)32 pub fn spawn_task(&mut self, task: impl Future<Output = ()> + 'a) { 33 let mut task = Box::pin(task); 34 // Schedule the task once. 35 match poll(&mut task.as_mut()) { 36 Some(_) => {} 37 _ => self.tasks.push(task), 38 } 39 } 40 41 /// Polls all `Future`s once. poll(&mut self)42 pub fn poll(&mut self) { 43 let mut idx = 0; 44 while let Some(task) = self.tasks.get_mut(idx) { 45 if poll(&mut task.as_mut()).is_some() { 46 let _ = self.tasks.swap_remove(idx); 47 } else { 48 idx += 1; 49 } 50 } 51 } 52 53 /// Runs all `Future`s until completion. run(&mut self)54 pub fn run(&mut self) { 55 while !self.tasks.is_empty() { 56 self.poll(); 57 } 58 } 59 60 /// Returns the current number of tasks. num_tasks(&self) -> usize61 pub fn num_tasks(&self) -> usize { 62 self.tasks.len() 63 } 64 } 65 66 #[cfg(test)] 67 mod test { 68 use super::*; 69 use gbl_async::yield_now; 70 use std::sync::Mutex; 71 72 #[test] test_spawn_and_poll_task()73 fn test_spawn_and_poll_task() { 74 let val1 = Mutex::new(0); 75 let val2 = Mutex::new(1); 76 77 let mut executor: CyclicExecutor = Default::default(); 78 // Spawns 2 tasks. 79 executor.spawn_task(async { 80 *val1.try_lock().unwrap() += 1; 81 yield_now().await; 82 *val1.try_lock().unwrap() += 1; 83 yield_now().await; 84 *val1.try_lock().unwrap() += 1; 85 yield_now().await; 86 }); 87 executor.spawn_task(async { 88 *val2.try_lock().unwrap() += 1; 89 yield_now().await; 90 *val2.try_lock().unwrap() += 1; 91 yield_now().await; 92 *val2.try_lock().unwrap() += 1; 93 yield_now().await; 94 }); 95 96 // Test that spawning a task schedules it immediately. 97 assert_eq!(*val1.try_lock().unwrap(), 1); 98 assert_eq!(*val2.try_lock().unwrap(), 2); 99 100 // Polls all Futures once. 101 executor.poll(); 102 assert_eq!(*val1.try_lock().unwrap(), 2); 103 assert_eq!(*val2.try_lock().unwrap(), 3); 104 105 // Runs to completion. 106 executor.run(); 107 assert_eq!(*val1.try_lock().unwrap(), 3); 108 assert_eq!(*val2.try_lock().unwrap(), 4); 109 } 110 111 #[test] test_complete_on_spawn_not_added()112 fn test_complete_on_spawn_not_added() { 113 let mut executor: CyclicExecutor = Default::default(); 114 executor.spawn_task(async {}); 115 assert_eq!(executor.num_tasks(), 0); 116 } 117 } 118