Skip to main content

strontium_core/
scheduler.rs

1use std::collections::{HashMap, HashSet, VecDeque};
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::{Arc, Mutex};
5use std::task::Wake;
6
7pub(crate) type TaskFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
8
9pub(crate) struct WakeQueue(pub(crate) Mutex<VecDeque<usize>>);
10
11impl WakeQueue {
12    pub(crate) fn new() -> Arc<Self> {
13        Arc::new(Self(Mutex::new(VecDeque::new())))
14    }
15
16    pub(crate) fn push(&self, id: usize) {
17        self.0.lock().expect("wake queue").push_back(id);
18    }
19}
20
21pub(crate) struct SimWaker {
22    pub(crate) task_id: usize,
23    pub(crate) queue: Arc<WakeQueue>,
24}
25
26impl Wake for SimWaker {
27    fn wake(self: Arc<Self>) {
28        self.queue.push(self.task_id);
29    }
30
31    fn wake_by_ref(self: &Arc<Self>) {
32        self.queue.push(self.task_id);
33    }
34}
35
36pub struct TaskHandle {
37    pub(crate) task_id: usize,
38    pub(crate) aborted: Arc<Mutex<HashSet<usize>>>,
39}
40
41impl TaskHandle {
42    pub fn abort(&self) {
43        self.aborted
44            .lock()
45            .expect("aborted set")
46            .insert(self.task_id);
47    }
48}
49
50pub(crate) struct Scheduler {
51    pub(crate) next_id: usize,
52    pub(crate) tasks: HashMap<usize, TaskFuture>,
53    pub(crate) pending_new: Vec<(usize, TaskFuture)>,
54    pub(crate) ready: VecDeque<usize>,
55    pub(crate) step_count: u64,
56    pub(crate) decision_log: Vec<usize>,
57    pub(crate) replay_decisions: VecDeque<usize>,
58}
59
60impl Scheduler {
61    pub(crate) fn new() -> Self {
62        Self {
63            next_id: 0,
64            tasks: HashMap::new(),
65            pending_new: Vec::new(),
66            ready: VecDeque::new(),
67            step_count: 0,
68            decision_log: Vec::new(),
69            replay_decisions: VecDeque::new(),
70        }
71    }
72}