Skip to main content

leodos_libcfs/runtime/
dyn_scope.rs

1//! A dynamic task scope for managing pinned futures.
2
3use core::future::Future;
4use core::pin::Pin;
5use core::task::{Context, Poll};
6use heapless::Vec;
7
8use crate::error::CfsError;
9
10/// A scope that holds Pinned mutable references to dyn-futures.
11pub struct DynScope<'a, const MAX_TASKS: usize = 8> {
12    tasks: Vec<Pin<&'a mut dyn Future<Output = Result<(), CfsError>>>, MAX_TASKS>,
13}
14
15impl<'a, const MAX_TASKS: usize> DynScope<'a, MAX_TASKS> {
16    /// Create a new DynScope
17    pub fn new() -> Self {
18        Self { tasks: Vec::new() }
19    }
20
21    /// Spawn a task into this scope.
22    ///
23    /// The argument must be a `Pin<&mut F>`. This ensures that the
24    /// underlying future has already been safely pinned by the caller.
25    pub fn spawn<F>(&mut self, future: Pin<&'a mut F>) -> Result<(), CfsError>
26    where
27        F: Future<Output = Result<(), CfsError>> + 'a,
28    {
29        self.tasks.push(future).map_err(|_| CfsError::TaskPoolFull)
30    }
31}
32
33impl<'a, const MAX_TASKS: usize> Future for DynScope<'a, MAX_TASKS> {
34    type Output = Result<(), CfsError>;
35
36    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
37        let tasks = &mut self.tasks;
38
39        let mut errors: Option<CfsError> = None;
40
41        let mut i = tasks.len();
42        while i > 0 {
43            i -= 1;
44
45            let pinned_task = &mut tasks[i];
46
47            match pinned_task.as_mut().poll(cx) {
48                Poll::Ready(Ok(())) => {
49                    tasks.swap_remove(i);
50                }
51                Poll::Ready(Err(e)) => {
52                    errors = Some(e);
53                    break;
54                }
55                Poll::Pending => {}
56            }
57        }
58
59        if let Some(e) = errors {
60            return Poll::Ready(Err(e));
61        }
62
63        if tasks.is_empty() {
64            Poll::Ready(Ok(()))
65        } else {
66            Poll::Pending
67        }
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use pin_utils::pin_mut;
75
76    #[test]
77    fn test_safe_scope() {
78        let s1 = async { Ok(()) };
79        let s2 = async { Ok(()) };
80        pin_mut!(s1);
81        pin_mut!(s2);
82        let mut scope = DynScope::<8>::new();
83        scope.spawn(s1).unwrap();
84        scope.spawn(s2).unwrap();
85    }
86}