tpt_archon_kernel/
scheduler.rs

1//! A cooperative async task scheduler (user-space first).
2//!
3//! One [`Task`] per database connection, not an OS process. This is a
4//! deterministic, single-threaded, cooperative round-robin scheduler suitable
5//! for the user-space validation model called for in `spec.txt`'s Risk 1
6//! mitigation (prove the architecture on a host OS before bare-metal / real
7//! `io_uring`).
8//!
9//! # Deadlock-freedom
10//!
11//! Because tasks are polled cooperatively and the scheduler holds no locks
12//! across `poll`, a task can only block progress by never yielding `Ready`.
13//! The scheduler always makes progress if any task is runnable and never waits
14//! on a task while holding a resource another task needs — the property
15//! `tpt-telos` is intended to prove (see `formal-proofs/`). Until then the
16//! behavior is exercised by the tests below.
17
18use alloc::boxed::Box;
19use alloc::collections::VecDeque;
20
21/// The result of polling a task once.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum Poll {
24    /// The task made progress but is not finished; re-schedule it.
25    Pending,
26    /// The task completed.
27    Ready,
28}
29
30/// A unit of schedulable work (e.g. one DB connection's driver).
31pub trait Task {
32    /// Advances the task. Returning [`Poll::Ready`] removes it from the
33    /// scheduler.
34    fn poll(&mut self) -> Poll;
35}
36
37/// A boxed task with an assigned id.
38struct Entry {
39    id: u64,
40    task: Box<dyn Task>,
41}
42
43/// A cooperative round-robin scheduler.
44#[derive(Default)]
45pub struct Scheduler {
46    ready: VecDeque<Entry>,
47    next_id: u64,
48}
49
50impl Scheduler {
51    /// Creates an empty scheduler.
52    pub fn new() -> Self {
53        Self {
54            ready: VecDeque::new(),
55            next_id: 0,
56        }
57    }
58
59    /// Spawns a task, returning its id.
60    pub fn spawn(&mut self, task: Box<dyn Task>) -> u64 {
61        let id = self.next_id;
62        self.next_id += 1;
63        self.ready.push_back(Entry { id, task });
64        id
65    }
66
67    /// Number of tasks still scheduled.
68    pub fn task_count(&self) -> usize {
69        self.ready.len()
70    }
71
72    /// Polls one task (round-robin). Returns the polled task's id and result,
73    /// or `None` if there are no tasks.
74    pub fn tick(&mut self) -> Option<(u64, Poll)> {
75        let mut entry = self.ready.pop_front()?;
76        let result = entry.task.poll();
77        let id = entry.id;
78        if result == Poll::Pending {
79            self.ready.push_back(entry);
80        }
81        Some((id, result))
82    }
83
84    /// Runs until every task completes. Returns the number of `tick`s executed.
85    ///
86    /// Guaranteed to terminate as long as every task eventually returns
87    /// [`Poll::Ready`].
88    pub fn run_to_completion(&mut self) -> usize {
89        let mut ticks = 0;
90        while self.tick().is_some() {
91            ticks += 1;
92        }
93        ticks
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use alloc::rc::Rc;
101    use core::cell::RefCell;
102
103    struct CountdownTask {
104        remaining: u32,
105        log: Rc<RefCell<alloc::vec::Vec<u64>>>,
106        id: u64,
107    }
108
109    impl Task for CountdownTask {
110        fn poll(&mut self) -> Poll {
111            self.log.borrow_mut().push(self.id);
112            if self.remaining == 0 {
113                Poll::Ready
114            } else {
115                self.remaining -= 1;
116                Poll::Pending
117            }
118        }
119    }
120
121    #[test]
122    fn runs_all_tasks_to_completion() {
123        let mut s = Scheduler::new();
124        let log = Rc::new(RefCell::new(alloc::vec::Vec::new()));
125        s.spawn(Box::new(CountdownTask {
126            remaining: 2,
127            log: log.clone(),
128            id: 0,
129        }));
130        s.spawn(Box::new(CountdownTask {
131            remaining: 1,
132            log: log.clone(),
133            id: 1,
134        }));
135        let ticks = s.run_to_completion();
136        assert_eq!(s.task_count(), 0);
137        // Round-robin interleaving proves fairness (no task starves).
138        assert!(ticks >= 5);
139        let l = log.borrow();
140        assert!(l.contains(&0) && l.contains(&1));
141    }
142
143    #[test]
144    fn immediate_ready_task_finishes_in_one_tick() {
145        struct Done;
146        impl Task for Done {
147            fn poll(&mut self) -> Poll {
148                Poll::Ready
149            }
150        }
151        let mut s = Scheduler::new();
152        let id = s.spawn(Box::new(Done));
153        assert_eq!(s.tick(), Some((id, Poll::Ready)));
154        assert_eq!(s.task_count(), 0);
155        assert_eq!(s.tick(), None);
156    }
157}
tpt_archon_kernel/scheduler.rs

tpt_archon_kernel/
scheduler.rs
