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

task/manager/
lifecycle.rs
