|
| 1 | +use super::{Task, TaskId}; |
| 2 | +use alloc::{collections::BTreeMap, sync::Arc, task::Wake}; |
| 3 | +use core::task::{Context, Poll, Waker}; |
| 4 | +use crossbeam_queue::ArrayQueue; |
| 5 | + |
| 6 | +pub struct Executor { |
| 7 | + tasks: BTreeMap<TaskId, Task>, |
| 8 | + task_queue: Arc<ArrayQueue<TaskId>>, |
| 9 | + waker_cache: BTreeMap<TaskId, Waker>, |
| 10 | +} |
| 11 | + |
| 12 | +impl Executor { |
| 13 | + pub fn new() -> Self { |
| 14 | + Executor { |
| 15 | + tasks: BTreeMap::new(), |
| 16 | + task_queue: Arc::new(ArrayQueue::new(100)), |
| 17 | + waker_cache: BTreeMap::new(), |
| 18 | + } |
| 19 | + } |
| 20 | + |
| 21 | + /// Spawns a new task. |
| 22 | + /// |
| 23 | + /// Because we are making a mutable loan from the executor, we can no longer execute `spawn` after the `run` |
| 24 | + /// method starts executing, plus `run` implements an infinite loop with a divergent return. |
| 25 | + /// One solution is to create a custom `Spawner` type that shares a queue with the `Executor`, a queue shared with it, |
| 26 | + /// or its own queue that is synchronized by the `Executor`. |
| 27 | + /// |
| 28 | + /// Remember that Rust doesn't allow having two mutable borrows at the same time, except for reborrowing. |
| 29 | + pub fn spawn(&mut self, task: Task) { |
| 30 | + let task_id = task.id; |
| 31 | + |
| 32 | + if self.tasks.insert(task.id, task).is_some() { |
| 33 | + panic!("task with same ID already in tasks"); |
| 34 | + } |
| 35 | + |
| 36 | + self.task_queue.push(task_id).expect("queue full"); |
| 37 | + } |
| 38 | + |
| 39 | + fn run_ready_tasks(&mut self) { |
| 40 | + // Destructuring is necessary because in the closure below we attempt to perform a full borrow of |
| 41 | + // self in order to obtain the waker_cache. |
| 42 | + let Self { |
| 43 | + tasks, |
| 44 | + task_queue, |
| 45 | + waker_cache, |
| 46 | + } = self; |
| 47 | + |
| 48 | + while let Some(task_id) = task_queue.pop() { |
| 49 | + let task = match tasks.get_mut(&task_id) { |
| 50 | + Some(task) => task, |
| 51 | + // Task no longer exists. |
| 52 | + None => continue, |
| 53 | + }; |
| 54 | + |
| 55 | + let waker = waker_cache |
| 56 | + .entry(task_id) |
| 57 | + .or_insert_with(|| TaskWaker::new(task_id, task_queue.clone())); |
| 58 | + |
| 59 | + let mut context = Context::from_waker(waker); |
| 60 | + |
| 61 | + match task.poll(&mut context) { |
| 62 | + Poll::Ready(()) => { |
| 63 | + // If the task is complete, remove it and its curly waker. There's no reason to keep them, |
| 64 | + // since the task is finished. |
| 65 | + tasks.remove(&task_id); |
| 66 | + waker_cache.remove(&task_id); |
| 67 | + } |
| 68 | + |
| 69 | + Poll::Pending => {} |
| 70 | + } |
| 71 | + } |
| 72 | + } |
| 73 | + |
| 74 | + /// The executor spins. |
| 75 | + /// |
| 76 | + /// Because the keyboard task, for example, prevents the tasks map from being empty, a loop with a |
| 77 | + /// divergent return value should model such tasks. |
| 78 | + pub fn run(&mut self) -> ! { |
| 79 | + loop { |
| 80 | + self.run_ready_tasks(); |
| 81 | + self.sleep_if_idle(); |
| 82 | + } |
| 83 | + } |
| 84 | + |
| 85 | + /// It puts the CPU into sleep mode when there are no tasks in the task queue, preventing the CPU from becoming busy. |
| 86 | + fn sleep_if_idle(&self) { |
| 87 | + use x86_64::instructions::interrupts::{self, enable_and_hlt}; |
| 88 | + |
| 89 | + interrupts::disable(); // Prevent race conditions |
| 90 | + // Between run_ready_tasks and sleep_if_idle, an interruption may occur and the queue may not become empty, hence the new check. |
| 91 | + if self.task_queue.is_empty() { |
| 92 | + // We disabled interrupts earlier because if an interrupt happens here, we'll lose the wakeup. |
| 93 | + // After verifying that there are indeed no tasks in the queue, we re-enable interrupts and activate |
| 94 | + // the hlt instruction to enter sleep mode. This is all done atomically. |
| 95 | + enable_and_hlt(); |
| 96 | + } else { |
| 97 | + // This means that after run_ready_tasks a new task was added by an interrupt, so we re-enable interrupts |
| 98 | + // and re-enter the loop. |
| 99 | + interrupts::enable(); |
| 100 | + } |
| 101 | + } |
| 102 | +} |
| 103 | + |
| 104 | +/// The waker's job is to push the waken task ID to the task_queue. |
| 105 | +/// Next, the `Executor` polls for the new task. |
| 106 | +struct TaskWaker { |
| 107 | + task_id: TaskId, |
| 108 | + // Ownership of task_queue is shared between wakers and executors through the Arc wrapper type, |
| 109 | + // which is based on reference counting. |
| 110 | + task_queue: Arc<ArrayQueue<TaskId>>, |
| 111 | +} |
| 112 | + |
| 113 | +impl TaskWaker { |
| 114 | + fn new(task_id: TaskId, task_queue: Arc<ArrayQueue<TaskId>>) -> Waker { |
| 115 | + // The Waker type supports conversions using the From trait when the type in question implements the Wake trait. |
| 116 | + // This is because we are wrapping a type that implements the Wake trait, where this trait uses the Arc smart pointer. |
| 117 | + Waker::from(Arc::new(TaskWaker { |
| 118 | + task_id, |
| 119 | + task_queue, |
| 120 | + })) |
| 121 | + } |
| 122 | + |
| 123 | + fn wake_task(&self) { |
| 124 | + self.task_queue.push(self.task_id).expect("task_queue full"); |
| 125 | + } |
| 126 | +} |
| 127 | + |
| 128 | +// We need to instantiate a Waker from the TaskWaker, where the simplest way is to |
| 129 | +// implement the Wake trait, which is based on Arc. |
| 130 | +// It's based on Arc because wakers are commonly shared between executors and asynchronous tasks. |
| 131 | +// The Waker type supports conversions using the From trait when the type in question implements the Wake trait. |
| 132 | +impl Wake for TaskWaker { |
| 133 | + // Since this captures ownership, it increases the number of references in Arc. |
| 134 | + fn wake(self: Arc<Self>) { |
| 135 | + self.wake_task(); |
| 136 | + } |
| 137 | + |
| 138 | + // Implementing this method is optional because not all data types support waking by reference. |
| 139 | + // However, implementing it provides performance benefits because it eliminates the need to modify |
| 140 | + // the reference count, for example. |
| 141 | + fn wake_by_ref(self: &Arc<Self>) { |
| 142 | + // Since our type only requires one &self reference, this is easy to resolve. |
| 143 | + self.wake_task(); |
| 144 | + } |
| 145 | +} |
0 commit comments