cortex_m_asyncrt/os/
executor.rs

1extern crate alloc;
2
3use super::*;
4use alloc::collections::BTreeMap;
5use alloc::sync::Arc;
6use alloc::task::Wake;
7use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
8use cortex_m::asm;
9use crossbeam_queue::ArrayQueue;
10
11struct TaskWaker {
12    task_id: TaskId,
13    task_queue: Arc<ArrayQueue<TaskId>>,
14}
15
16impl TaskWaker {
17    fn new(task_id: TaskId, task_queue: Arc<ArrayQueue<TaskId>>) -> Waker {
18        Waker::from(Arc::new(TaskWaker {
19            task_id,
20            task_queue,
21        }))
22    }
23    fn wake_task(&self) {
24        self.task_queue
25            .push(self.task_id)
26            .expect("Task queue is full!");
27    }
28}
29
30impl Wake for TaskWaker {
31    fn wake(self: Arc<Self>) {
32        self.wake_task();
33    }
34
35    fn wake_by_ref(self: &Arc<Self>) {
36        self.wake_task();
37    }
38}
39
40/// Task executor that runs tasks to completion.
41pub struct Executor {
42    /// The tasks to run are stored in a BTreeMap. This allows us to store tasks
43    /// in a way that they can be accessed by their ID. BtreeMap is used because
44    /// it is a sorted map and we can use the task ID as the key.
45    tasks: BTreeMap<TaskId, Task>,
46    /// The task queue is a queue of task IDs. This queue is used to keep track of
47    /// the order in which tasks are spawned. The executor will run tasks in the
48    /// order they were spawned.
49    task_queue: Arc<ArrayQueue<TaskId>>,
50    /// The waker cache is a map of task IDs to wakers. This cache is used to store
51    /// wakers for tasks that are currently running. This allows us to wake up tasks
52    /// when they are ready to be polled again.
53    waker_cache: BTreeMap<TaskId, Waker>,
54}
55
56impl Executor {
57    /// Create a new executor with an empty task queue.
58    pub fn new<const N: usize>() -> Executor {
59        Executor {
60            tasks: BTreeMap::new(),
61            task_queue: Arc::new(ArrayQueue::new(N)),
62            waker_cache: BTreeMap::new(),
63        }
64    }
65
66    /// Spawn a new task. This function pushes the task to the back of the task queue.
67    pub fn spawn(&mut self, task: Task) {
68        let task_id = task.id;
69        if self.tasks.insert(task_id, task).is_some() {
70            panic!("task with same ID already in tasks");
71        }
72        self.task_queue.push(task_id).expect("Task queue is full.");
73    }
74
75    /// Run the executor. This function runs tasks to completion.
76    fn run_ready_tasks(&mut self) {
77        // Use a while let loop to pop task_ids from the task queue.
78        while let Some(task_id) = self.task_queue.pop() {
79            // Get the task from the tasks map.
80            let task = match self.tasks.get_mut(&task_id) {
81                Some(task) => task,
82                None => continue, // task was removed from the tasks map so we skip it.
83            };
84
85            // Get the waker for the task from the waker cache.
86            let waker = self
87                .waker_cache
88                .entry(task_id)
89                .or_insert_with(|| TaskWaker::new(task_id, self.task_queue.clone())); // create a new waker if it doesn't exist.
90
91            // Create a new context from the waker.
92            let mut context = Context::from_waker(waker);
93
94            // Poll the task.
95            match task.poll(&mut context) {
96                Poll::Ready(()) => {
97                    // task is done, so we remove it and its cached waker.
98                    self.tasks.remove(&task_id);
99                    self.waker_cache.remove(&task_id);
100                }
101                Poll::Pending => {} // task is not done, so we do nothing.
102            }
103        }
104    }
105
106    /// Put the processor to sleep if there are no tasks to run.
107    fn sleep_if_idle(&self) {
108        // If there are no tasks to run, we put the processor to sleep until an interrupt occurs.
109        cortex_m::interrupt::free(|_| {
110            if self.task_queue.is_empty() {
111                asm::wfi(); // Wait for interrupt.
112            }
113        });
114    }
115
116    /// Run the executor. This function will run tasks until the task queue is empty.
117    pub fn run(&mut self) {
118        loop {
119            self.run_ready_tasks();
120            self.sleep_if_idle();
121        }
122    }
123}
124
125pub fn dummy_raw_waker() -> RawWaker {
126    fn no_op(_: *const ()) {}
127    fn clone(_: *const ()) -> RawWaker {
128        dummy_raw_waker()
129    }
130
131    let vtable = &RawWakerVTable::new(clone, no_op, no_op, no_op);
132
133    RawWaker::new(0 as *const (), vtable)
134}
135
136pub fn dummy_waker() -> Waker {
137    unsafe { Waker::from_raw(dummy_raw_waker()) }
138}