1use std::collections::HashMap;
20
21pub type Ms = u64;
23
24const MAX_STRETCH: u32 = 4;
26const CALM_RUNS: u32 = 3;
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
31pub enum SourceId {
32 Machine,
34 Panes,
37}
38
39impl SourceId {
40 pub fn name(self) -> &'static str {
41 match self {
42 Self::Machine => "machine",
43 Self::Panes => "panes",
44 }
45 }
46}
47
48#[derive(Debug, Clone)]
50pub struct SourceSpec {
51 pub id: SourceId,
52 pub min_interval_ms: Ms,
54 pub budget_ms: Ms,
56 pub needs: &'static [SourceId],
59}
60
61#[derive(Debug, Clone)]
62struct Lease {
63 holder: String,
64 source: SourceId,
65 interval_ms: Ms,
66 expires_at: Option<Ms>,
68}
69
70#[derive(Debug, Clone, Default)]
71struct SourceState {
72 next_due: Option<Ms>,
74 last_start: Option<Ms>,
75 last_took: Ms,
76 stretch: u32,
78 calm: u32,
79 runs: u64,
80 overruns: u64,
81 skipped: u64,
83}
84
85#[derive(Debug, Clone, PartialEq, Eq)]
87pub struct SourceStatus {
88 pub id: SourceId,
89 pub holders: Vec<String>,
90 pub interval_ms: Option<Ms>,
92 pub stretch: u32,
93 pub last_took_ms: Ms,
94 pub runs: u64,
95 pub overruns: u64,
96 pub skipped: u64,
97}
98
99pub struct Scheduler {
100 specs: Vec<SourceSpec>,
101 leases: Vec<Lease>,
102 state: HashMap<SourceId, SourceState>,
103 coalesce_ms: Ms,
106}
107
108impl Scheduler {
109 pub fn new(specs: Vec<SourceSpec>, coalesce_ms: Ms) -> Self {
111 for (i, s) in specs.iter().enumerate() {
112 assert!(
113 s.needs.iter().all(|n| specs[..i].iter().any(|p| p.id == *n)),
114 "source {:?} needs a source declared after it",
115 s.id
116 );
117 }
118 let state = specs.iter().map(|s| (s.id, SourceState { stretch: 1, ..Default::default() })).collect();
119 Self { specs, leases: Vec::new(), state, coalesce_ms }
120 }
121
122 pub fn lease(&mut self, holder: &str, source: SourceId, interval_ms: Ms, expires_at: Option<Ms>, now: Ms) {
126 match self.leases.iter_mut().find(|l| l.holder == holder && l.source == source) {
127 Some(l) => {
128 l.interval_ms = interval_ms;
129 l.expires_at = expires_at;
130 }
131 None => self.leases.push(Lease { holder: holder.to_string(), source, interval_ms, expires_at }),
132 }
133 self.reschedule(source, now);
134 }
135
136 pub fn release(&mut self, holder: &str, now: Ms) {
138 let sources: Vec<SourceId> = self.leases.iter().filter(|l| l.holder == holder).map(|l| l.source).collect();
139 self.leases.retain(|l| l.holder != holder);
140 for source in sources {
141 self.reschedule(source, now);
142 }
143 }
144
145 pub fn release_source(&mut self, holder: &str, source: SourceId, now: Ms) {
147 self.leases.retain(|l| !(l.holder == holder && l.source == source));
148 self.reschedule(source, now);
149 }
150
151 fn reschedule(&mut self, source: SourceId, now: Ms) {
155 let Some(interval) = self.effective_interval_ms(source, now) else {
156 if let Some(st) = self.state.get_mut(&source) {
159 st.next_due = None;
160 }
161 return;
162 };
163 if let Some(st) = self.state.get_mut(&source) {
164 if let (Some(_), Some(last)) = (st.next_due, st.last_start) {
165 st.next_due = Some((last + interval).max(last + st.last_took));
166 }
167 }
168 }
169
170 fn spec(&self, id: SourceId) -> &SourceSpec {
171 self.specs.iter().find(|s| s.id == id).expect("every SourceId has a spec")
172 }
173
174 fn live(&self, source: SourceId, now: Ms) -> impl Iterator<Item = &Lease> {
175 self.leases.iter().filter(move |l| l.source == source && l.expires_at.is_none_or(|e| e > now))
176 }
177
178 pub fn effective_interval_ms(&self, source: SourceId, now: Ms) -> Option<Ms> {
181 let wanted = self.live(source, now).map(|l| l.interval_ms).min()?;
182 let stretch = self.state.get(&source).map_or(1, |s| s.stretch.max(1));
183 Some(wanted.max(self.spec(source).min_interval_ms) * Ms::from(stretch))
184 }
185
186 pub fn due(&mut self, now: Ms) -> Vec<SourceId> {
189 self.leases.retain(|l| l.expires_at.is_none_or(|e| e > now));
190 let ids: Vec<SourceId> = self.specs.iter().map(|s| s.id).collect();
191 let mut run = Vec::new();
192 for id in ids {
193 let leased = self.live(id, now).next().is_some();
194 let st = self.state.get_mut(&id).expect("state for every source");
195 if !leased {
196 st.next_due = None;
197 continue;
198 }
199 let due = *st.next_due.get_or_insert(now);
200 if due <= now + self.coalesce_ms {
201 run.push(id);
202 }
203 }
204 let mut wanted = run.clone();
206 for id in &run {
207 wanted.extend_from_slice(self.spec(*id).needs);
208 }
209 self.specs.iter().map(|s| s.id).filter(|id| wanted.contains(id)).collect()
210 }
211
212 pub fn next_wakeup(&self, now: Ms) -> Option<Ms> {
214 self.specs
215 .iter()
216 .filter(|s| self.live(s.id, now).next().is_some())
217 .map(|s| self.state[&s.id].next_due.unwrap_or(now).max(now))
218 .min()
219 }
220
221 pub fn finished(&mut self, id: SourceId, started: Ms, took_ms: Ms) {
225 let budget = self.spec(id).budget_ms;
226 let interval_before = self.effective_interval_ms(id, started).unwrap_or(1).max(1);
227 let st = self.state.get_mut(&id).expect("state for every source");
228 let planned = st.next_due.unwrap_or(started);
229 st.skipped += started.saturating_sub(planned) / interval_before;
230 st.runs += 1;
231 st.last_start = Some(started);
232 st.last_took = took_ms;
233 if took_ms > budget {
234 st.overruns += 1;
235 st.calm = 0;
236 st.stretch = (st.stretch.max(1) * 2).min(MAX_STRETCH);
237 } else if took_ms * 2 <= budget {
238 st.calm += 1;
239 if st.calm >= CALM_RUNS && st.stretch > 1 {
240 st.stretch /= 2;
241 st.calm = 0;
242 }
243 } else {
244 st.calm = 0;
245 }
246 let end = started + took_ms;
247 let interval = self.effective_interval_ms(id, end).unwrap_or(interval_before);
248 let st = self.state.get_mut(&id).expect("state for every source");
249 st.next_due = Some((started + interval).max(end));
250 }
251
252 pub fn status(&self, now: Ms) -> Vec<SourceStatus> {
253 self.specs
254 .iter()
255 .map(|s| {
256 let st = &self.state[&s.id];
257 SourceStatus {
258 id: s.id,
259 holders: self.live(s.id, now).map(|l| l.holder.clone()).collect(),
260 interval_ms: self.effective_interval_ms(s.id, now),
261 stretch: st.stretch.max(1),
262 last_took_ms: st.last_took,
263 runs: st.runs,
264 overruns: st.overruns,
265 skipped: st.skipped,
266 }
267 })
268 .collect()
269 }
270}
271
272#[cfg(test)]
273mod tests {
274 use super::*;
275 use SourceId::{Machine, Panes};
276
277 fn spec(id: SourceId, min: Ms, budget: Ms, needs: &'static [SourceId]) -> SourceSpec {
278 SourceSpec { id, min_interval_ms: min, budget_ms: budget, needs }
279 }
280
281 fn sched() -> Scheduler {
282 Scheduler::new(vec![spec(Machine, 200, 100, &[]), spec(Panes, 500, 200, &[])], 50)
283 }
284
285 fn tick(s: &mut Scheduler, now: Ms, took: Ms) -> Vec<SourceId> {
287 let ran = s.due(now);
288 for id in &ran {
289 s.finished(*id, now, took);
290 }
291 ran
292 }
293
294 #[test]
295 fn nothing_leased_means_nothing_due_and_no_wakeup() {
296 let mut s = sched();
297 assert!(s.due(0).is_empty());
298 assert_eq!(s.next_wakeup(0), None);
299 }
300
301 #[test]
302 fn a_new_lease_runs_at_once_then_every_interval() {
303 let mut s = sched();
304 s.lease("bar", Machine, 1000, None, 0);
305 assert_eq!(tick(&mut s, 0, 10), [Machine]);
306 assert_eq!(s.next_wakeup(0), Some(1000));
307 assert!(tick(&mut s, 500, 10).is_empty());
308 assert_eq!(tick(&mut s, 1000, 10), [Machine]);
309 assert_eq!(s.next_wakeup(1000), Some(2000));
310 }
311
312 #[test]
313 fn the_shortest_live_lease_wins_and_releasing_it_slows_the_source() {
314 let mut s = sched();
315 s.lease("janitor", Panes, 30_000, None, 0);
316 tick(&mut s, 0, 10);
317 assert_eq!(s.effective_interval_ms(Panes, 0), Some(30_000));
318 assert_eq!(s.next_wakeup(0), Some(30_000));
319 s.lease("badge", Panes, 1000, None, 100);
321 assert_eq!(s.effective_interval_ms(Panes, 100), Some(1000));
322 assert_eq!(s.next_wakeup(100), Some(1000));
323 tick(&mut s, 1000, 10);
324 s.release("badge", 1100);
325 assert_eq!(s.effective_interval_ms(Panes, 1100), Some(30_000));
326 assert_eq!(s.next_wakeup(1100), Some(31_000));
328 assert!(tick(&mut s, 2000, 10).is_empty());
329 }
330
331 #[test]
332 fn releasing_the_last_lease_stops_the_source() {
333 let mut s = sched();
334 s.lease("badge", Panes, 1000, None, 0);
335 tick(&mut s, 0, 10);
336 s.release_source("badge", Panes, 500);
337 assert_eq!(s.next_wakeup(500), None);
338 assert!(tick(&mut s, 1000, 10).is_empty());
339 }
340
341 #[test]
342 fn a_lease_taken_right_after_the_last_was_released_starts_the_source_at_once() {
343 let mut s = sched();
344 s.lease("badge", Panes, 30_000, None, 0);
345 tick(&mut s, 0, 10);
346 s.release_source("badge", Panes, 500);
348 s.lease("badge", Panes, 30_000, None, 600);
349 assert_eq!(s.next_wakeup(600), Some(600));
350 assert_eq!(tick(&mut s, 600, 10), [Panes]);
351 }
352
353 #[test]
354 fn an_expired_lease_is_forgotten_the_same_way_when_a_new_one_follows() {
355 let mut s = sched();
356 s.lease("badge", Panes, 30_000, Some(1000), 0);
357 tick(&mut s, 0, 10);
358 s.lease("other", Panes, 30_000, None, 2000);
359 assert_eq!(s.next_wakeup(2000), Some(30_000), "the other lease is the only one, and it just ran at 0 on a 30 s interval");
360 }
361
362 #[test]
363 fn renewing_a_lease_with_a_longer_interval_pushes_the_next_run_out() {
364 let mut s = sched();
365 s.lease("bar", Machine, 1000, None, 0);
366 tick(&mut s, 0, 10);
367 s.lease("bar", Machine, 5000, None, 200);
368 assert_eq!(s.next_wakeup(200), Some(5000));
369 }
370
371 #[test]
372 fn a_source_never_runs_faster_than_its_own_minimum() {
373 let mut s = sched();
374 s.lease("fast", Machine, 10, None, 0);
375 assert_eq!(s.effective_interval_ms(Machine, 0), Some(200));
376 }
377
378 #[test]
379 fn an_expired_lease_stops_the_source_and_a_new_one_restarts_it_at_once() {
380 let mut s = sched();
381 s.lease("popover", Machine, 1000, Some(3000), 0);
382 assert_eq!(tick(&mut s, 0, 10), [Machine]);
383 assert_eq!(tick(&mut s, 2000, 10), [Machine]);
384 assert!(tick(&mut s, 3000, 10).is_empty(), "expired exactly at 3000");
385 assert_eq!(s.next_wakeup(3000), None);
386 s.lease("popover", Machine, 1000, Some(9000), 5000);
387 assert_eq!(tick(&mut s, 5000, 10), [Machine], "no waiting out the old schedule");
388 }
389
390 #[test]
391 fn renewing_a_lease_replaces_it() {
392 let mut s = sched();
393 s.lease("a", Machine, 1000, Some(2000), 0);
394 s.lease("a", Machine, 1000, Some(8000), 1500);
395 assert_eq!(s.status(1500)[0].holders, ["a"]);
396 assert_eq!(tick(&mut s, 3000, 10), [Machine]);
397 }
398
399 #[test]
400 fn sources_due_close_together_run_in_one_pass() {
401 let mut s = sched();
402 s.lease("a", Machine, 1000, None, 0);
403 s.lease("a", Panes, 1000, None, 0);
404 tick(&mut s, 0, 10);
405 tick(&mut s, 0, 10);
406 s.finished(Panes, 30, 10);
409 assert_eq!(s.due(1000), [Machine, Panes]);
410 }
411
412 #[test]
413 fn what_a_due_source_reads_runs_with_it_even_without_a_lease() {
414 let mut s = Scheduler::new(vec![spec(Machine, 200, 100, &[]), spec(Panes, 500, 200, &[Machine])], 50);
415 s.lease("badge", Panes, 1000, None, 0);
416 assert_eq!(s.due(0), [Machine, Panes]);
417 assert!(s.status(0)[0].holders.is_empty());
419 }
420
421 #[test]
422 #[should_panic(expected = "declared after")]
423 fn a_need_declared_later_is_refused() {
424 Scheduler::new(vec![spec(Panes, 500, 200, &[Machine]), spec(Machine, 200, 100, &[])], 50);
425 }
426
427 #[test]
428 fn a_source_over_budget_is_stretched_up_to_four_times_and_eases_back_when_calm() {
429 let mut s = sched();
430 s.lease("bar", Panes, 1000, None, 0);
431 let mut now = 0;
432 for expected in [2, 4, 4] {
433 tick(&mut s, now, 400); assert_eq!(s.status(now)[1].stretch, expected);
435 now = s.next_wakeup(now).unwrap();
436 }
437 assert_eq!(s.status(now)[1].overruns, 3);
438 assert_eq!(s.effective_interval_ms(Panes, now), Some(4000));
439 for _ in 0..3 {
441 tick(&mut s, now, 50);
442 now = s.next_wakeup(now).unwrap();
443 }
444 assert_eq!(s.status(now)[1].stretch, 2);
445 for _ in 0..3 {
446 tick(&mut s, now, 50);
447 now = s.next_wakeup(now).unwrap();
448 }
449 assert_eq!(s.status(now)[1].stretch, 1);
450 }
451
452 #[test]
453 fn a_run_in_the_middle_of_the_budget_neither_stretches_nor_eases() {
454 let mut s = sched();
455 s.lease("bar", Panes, 1000, None, 0);
456 tick(&mut s, 0, 400);
457 assert_eq!(s.status(0)[1].stretch, 2);
458 let mut now = s.next_wakeup(0).unwrap();
459 for _ in 0..6 {
460 tick(&mut s, now, 150); now = s.next_wakeup(now).unwrap();
462 }
463 assert_eq!(s.status(now)[1].stretch, 2);
464 }
465
466 #[test]
467 fn a_late_tick_is_counted_and_never_made_up_for() {
468 let mut s = sched();
469 s.lease("bar", Machine, 1000, None, 0);
470 tick(&mut s, 0, 10);
471 assert_eq!(tick(&mut s, 4500, 10), [Machine]);
473 assert_eq!(s.status(4500)[0].skipped, 3);
474 assert_eq!(s.next_wakeup(4500), Some(5500), "one interval on, not four runs in a burst");
475 }
476
477 #[test]
478 fn a_run_longer_than_its_interval_is_not_started_again_before_it_ends() {
479 let mut s = sched();
480 s.lease("bar", Machine, 200, None, 0);
481 tick(&mut s, 0, 900);
482 assert!(s.next_wakeup(0).unwrap() >= 900);
483 }
484
485 #[test]
486 fn status_names_holders_interval_and_counts() {
487 let mut s = sched();
488 s.lease("bar", Machine, 1000, None, 0);
489 s.lease("popover", Machine, 2000, None, 0);
490 tick(&mut s, 0, 10);
491 let m = &s.status(0)[0];
492 assert_eq!(m.id, Machine);
493 assert_eq!(m.holders, ["bar", "popover"]);
494 assert_eq!((m.interval_ms, m.runs, m.last_took_ms), (Some(1000), 1, 10));
495 let p = &s.status(0)[1];
496 assert_eq!((p.id, p.interval_ms, p.runs), (Panes, None, 0));
497 }
498}