xref: /aosp_15_r20/bootable/libbootloader/gbl/libasync/src/cyclic_executor.rs (revision 5225e6b173e52d2efc6bcf950c27374fd72adabc)
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