task/manager/
lifecycle.rs

1// Lifecycle module - handles task spawning, execution, and lifecycle management
2
3use super::*;
4use alloc::boxed::Box;
5use core::{
6    future::{Future, poll_fn},
7    ptr::NonNull,
8    task::Poll,
9    time::Duration,
10};
11use embassy_executor::raw::{TaskPool, task_from_waker};
12use embassy_time::Timer;
13
14impl Manager {
15    // Static function to create and execute tasks
16    // This function is outside the closure that captures SpawnToken,
17    // so it can be called safely from nested tasks
18    async fn create_and_run_task<R: 'static, FunctionType, FutureType>(
19        manager: &'static Manager,
20        parent_task_identifier: TaskIdentifier,
21        name: &str,
22        function: FunctionType,
23        spawner: Option<usize>,
24    ) -> Result<(JoinHandle<R>, TaskIdentifier)>
25    where
26        FunctionType: FnOnce(TaskIdentifier) -> FutureType + 'static,
27        FutureType: Future<Output = R> + 'static,
28    {
29        let identifier = manager
30            .register(parent_task_identifier, name)
31            .await
32            .expect("Failed to get new task identifier");
33
34        let pool = Box::new(TaskPool::<_, 1>::new());
35        let pool = Box::leak(pool);
36
37        let (join_handle_parent, join_handle_child) = JoinHandle::new();
38
39        let task = async move || {
40            let manager = get_instance();
41
42            let internal_identifier = Manager::get_current_internal_identifier().await;
43
44            manager
45                .set_internal_identifier(identifier, internal_identifier)
46                .await
47                .expect("Failed to register task");
48
49            let result = function(identifier).await;
50
51            join_handle_child.signal(result);
52
53            manager
54                .unregister(identifier)
55                .await
56                .expect("Failed to unregister task");
57        };
58
59        let mut inner = manager.0.write().await;
60
61        // Select the best spawner for the new task
62        let spawner = if let Some(spawner) = spawner {
63            if !inner.spawners.contains_key(&spawner) {
64                return Err(Error::InvalidSpawnerIdentifier);
65            }
66            spawner
67        } else {
68            Manager::select_best_spawner(&inner)?
69        };
70
71        inner
72            .tasks
73            .get_mut(&identifier)
74            .expect("Failed to get task metadata")
75            .spawner_identifier = spawner;
76
77        let token = pool.spawn(task).expect("Failed to spawn task");
78
79        inner
80            .spawners
81            .get(&spawner)
82            .expect("Failed to get spawner")
83            .spawn(token);
84
85        Ok((join_handle_parent, identifier))
86    }
87
88    /// Spawn task
89    pub async fn spawn<FunctionType, FutureType, ReturnType>(
90        &'static self,
91        parent_task: TaskIdentifier,
92        name: &str,
93        spawner: Option<usize>,
94        function: FunctionType,
95    ) -> Result<(JoinHandle<ReturnType>, TaskIdentifier)>
96    where
97        FunctionType: FnOnce(TaskIdentifier) -> FutureType + 'static,
98        FutureType: Future<Output = ReturnType> + 'static,
99        ReturnType: 'static,
100    {
101        // Call the helper function with all our parameters
102        Self::create_and_run_task(self, parent_task, name, function, spawner).await
103    }
104
105    /// Set the internal identifier of a task.
106    ///
107    /// This function check if the task identifier is not already used,
108    /// however it doesn't check if the parent task exists.
109    async fn set_internal_identifier(
110        &self,
111        identifier: TaskIdentifier,
112        internal_identifier: usize,
113    ) -> Result<()> {
114        let mut inner = self.0.write().await;
115
116        let metadata = Self::get_task_mutable(&mut inner, identifier)?;
117
118        metadata.internal_identifier = internal_identifier;
119
120        // Register the internal identifier of the task
121        if let Some(old_identifier) = inner.identifiers.insert(internal_identifier, identifier) {
122            // Rollback the task registration if internal identifier registration fails
123            inner.identifiers.remove(&internal_identifier);
124            inner
125                .identifiers
126                .insert(internal_identifier, old_identifier);
127            return Err(Error::InvalidTaskIdentifier);
128        }
129
130        Ok(())
131    }
132
133    /// Sleep the current thread for a given duration.
134    pub async fn sleep(duration: Duration) {
135        let nano_seconds = duration.as_nanos();
136
137        Timer::after(embassy_time::Duration::from_nanos(nano_seconds as u64)).await
138    }
139
140    pub async fn get_current_internal_identifier() -> usize {
141        poll_fn(|context| {
142            let task_reference = task_from_waker(context.waker());
143
144            let inner: NonNull<u8> = unsafe { core::mem::transmute(task_reference) };
145
146            let identifier = inner.as_ptr() as usize;
147
148            Poll::Ready(identifier)
149        })
150        .await
151    }
152
153    pub async fn get_current_task_identifier(&self) -> TaskIdentifier {
154        let internal_identifier = Self::get_current_internal_identifier().await;
155
156        *self
157            .0
158            .read()
159            .await
160            .identifiers
161            .get(&internal_identifier)
162            .expect("Failed to get task identifier")
163    }
164}
task/manager/lifecycle.rs

task/manager/
lifecycle.rs
