rust-kernel / src / task / executor.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
//! The real executor: only polls tasks that were *woken*.
//!
//! * `tasks`:      all live tasks, by id
//! * `task_queue`: ids of tasks ready to run. Wakers push onto it -- possibly
//!                 from an interrupt handler, so it is a lock-free queue and
//!                 pushing never allocates.
//! * `waker_cache`: one Waker per task, created once and reused.
//!
//! When nothing is ready, the CPU sleeps with `hlt` until the next interrupt.

use super::{Task, TaskId};
use alloc::{collections::BTreeMap, sync::Arc, task::Wake};
use core::task::{Context, Poll, Waker};
use crossbeam_queue::ArrayQueue;

pub struct Executor {
    tasks: BTreeMap<TaskId, Task>,
    task_queue: Arc<ArrayQueue<TaskId>>,
    waker_cache: BTreeMap<TaskId, Waker>,
}

impl Executor {
    pub fn new() -> Self {
        Executor {
            tasks: BTreeMap::new(),
            task_queue: Arc::new(ArrayQueue::new(100)),
            waker_cache: BTreeMap::new(),
        }
    }

    pub fn spawn(&mut self, task: Task) {
        let task_id = task.id;
        if self.tasks.insert(task.id, task).is_some() {
            panic!("task with same ID already in tasks");
        }
        self.task_queue.push(task_id).expect("queue full");
    }

    /// Run tasks forever (the kernel's main loop).
    pub fn run(&mut self) -> ! {
        loop {
            self.run_ready_tasks();
            self.sleep_if_idle();
        }
    }

    /// Run until every task has finished (used by tests).
    pub fn run_until_done(&mut self) {
        while !self.tasks.is_empty() {
            self.run_ready_tasks();
            self.sleep_if_idle();
        }
    }

    fn run_ready_tasks(&mut self) {
        // Destructure `self` so we can borrow the fields separately.
        let Self { tasks, task_queue, waker_cache } = self;

        while let Some(task_id) = task_queue.pop() {
            let task = match tasks.get_mut(&task_id) {
                Some(task) => task,
                None => continue, // task already finished (stale wake-up)
            };
            let waker = waker_cache
                .entry(task_id)
                .or_insert_with(|| TaskWaker::new(task_id, task_queue.clone()));
            let mut context = Context::from_waker(waker);
            if let Poll::Ready(()) = task.poll(&mut context) {
                tasks.remove(&task_id);
                waker_cache.remove(&task_id);
            }
        }
    }

    /// `hlt` until the next interrupt, but only if the queue is still empty.
    /// Interrupts are disabled for the check: otherwise an interrupt could
    /// wake a task *between* the check and `hlt`, and we'd sleep with work
    /// waiting. `enable_and_hlt` is `sti; hlt`, which the CPU executes
    /// atomically (interrupts are only recognised after the `hlt` starts).
    fn sleep_if_idle(&self) {
        use x86_64::instructions::interrupts::{self, enable_and_hlt};
        if crate::thread::scheduler_started() {
            // Block this thread until a waker fires (e.g. from the keyboard
            // interrupt): other threads get the CPU meanwhile.
            EXECUTOR_IDLE.wait_until(|| (!self.task_queue.is_empty()).then_some(()));
            return;
        }
        // No scheduler (some tests): just halt until the next interrupt.
        interrupts::disable();
        if self.task_queue.is_empty() {
            enable_and_hlt();
        } else {
            interrupts::enable();
        }
    }
}

impl Default for Executor {
    fn default() -> Self {
        Self::new()
    }
}

/// Waking = pushing the task id back onto the ready queue.
struct TaskWaker {
    task_id: TaskId,
    task_queue: Arc<ArrayQueue<TaskId>>,
}

impl TaskWaker {
    fn new(task_id: TaskId, task_queue: Arc<ArrayQueue<TaskId>>) -> Waker {
        // `Wake` trait + Arc -> Waker conversion from the alloc crate.
        Waker::from(Arc::new(TaskWaker { task_id, task_queue }))
    }

    fn wake_task(&self) {
        self.task_queue.push(self.task_id).expect("task_queue full");
        EXECUTOR_IDLE.wake_all();
    }
}

/// Where an idle executor thread blocks until a task is woken.
static EXECUTOR_IDLE: crate::thread::WaitQueue = crate::thread::WaitQueue::new();

impl Wake for TaskWaker {
    fn wake(self: Arc<Self>) {
        self.wake_task();
    }
    fn wake_by_ref(self: &Arc<Self>) {
        self.wake_task();
    }
}