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

task/manager/
lifecycle.rs
