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);
79
80        inner
81            .spawners
82            .get(&spawner)
83            .expect("Failed to get spawner")
84            .spawn(token)
85            .expect("Failed to spawn task");
86
87        Ok((join_handle_parent, identifier))
88    }
89
90    /// Spawn task
91    pub async fn spawn<FunctionType, FutureType, ReturnType>(
92        &'static self,
93        parent_task: TaskIdentifier,
94        name: &str,
95        spawner: Option<usize>,
96        function: FunctionType,
97    ) -> Result<(JoinHandle<ReturnType>, TaskIdentifier)>
98    where
99        FunctionType: FnOnce(TaskIdentifier) -> FutureType + 'static,
100        FutureType: Future<Output = ReturnType> + 'static,
101        ReturnType: 'static,
102    {
103        // Call the helper function with all our parameters
104        Self::create_and_run_task(self, parent_task, name, function, spawner).await
105    }
106
107    /// Set the internal identifier of a task.
108    ///
109    /// This function check if the task identifier is not already used,
110    /// however it doesn't check if the parent task exists.
111    async fn set_internal_identifier(
112        &self,
113        identifier: TaskIdentifier,
114        internal_identifier: usize,
115    ) -> Result<()> {
116        let mut inner = self.0.write().await;
117
118        let metadata = Self::get_task_mutable(&mut inner, identifier)?;
119
120        metadata.internal_identifier = internal_identifier;
121
122        // Register the internal identifier of the task
123        if let Some(old_identifier) = inner.identifiers.insert(internal_identifier, identifier) {
124            // Rollback the task registration if internal identifier registration fails
125            inner.identifiers.remove(&internal_identifier);
126            inner
127                .identifiers
128                .insert(internal_identifier, old_identifier);
129            return Err(Error::InvalidTaskIdentifier);
130        }
131
132        Ok(())
133    }
134
135    pub async fn r#yield() {
136        yield_now().await;
137    }
138
139    /// Sleep the current thread for a given duration.
140    pub async fn sleep(duration: Duration) {
141        let nano_seconds = duration.as_nanos();
142
143        Timer::after(embassy_time::Duration::from_nanos(nano_seconds as u64)).await
144    }
145
146    pub async fn get_current_internal_identifier() -> usize {
147        poll_fn(|context| {
148            let task_reference = task_from_waker(context.waker());
149
150            let inner: NonNull<u8> = unsafe { core::mem::transmute(task_reference) };
151
152            let identifier = inner.as_ptr() as usize;
153
154            Poll::Ready(identifier)
155        })
156        .await
157    }
158
159    pub async fn get_current_task_identifier(&self) -> TaskIdentifier {
160        let internal_identifier = Self::get_current_internal_identifier().await;
161
162        *self
163            .0
164            .read()
165            .await
166            .identifiers
167            .get(&internal_identifier)
168            .expect("Failed to get task identifier")
169    }
170}
task/manager/lifecycle.rs

task/manager/
lifecycle.rs
