Skip to main content

agentmux_srv/backend/sampler/
schedule.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Which sources to sample now, and how soon to look again.
5//!
6//! A **source** is one thing worth measuring (the machine's counters, the
7//! processes of each pane). Nothing runs unless someone holds a **lease** on it:
8//! a named holder asking for it at an interval, until it is released or its
9//! expiry passes. A source runs at the shortest interval among its live leases,
10//! never faster than its own minimum, and slows down while it keeps taking longer
11//! than its budget.
12//!
13//! This is the whole decision, kept pure so it is tested with a fake clock:
14//! [`Scheduler::due`] says what to run, the caller runs it and reports
15//! [`Scheduler::finished`], [`Scheduler::next_wakeup`] says when to look again.
16//! Times are milliseconds on one monotonic clock the caller owns.
17//! docs/reports/REPORT_SYSINFO_REWRITE_AND_MODERN_MONITOR_2026_10_13.md section 5.5.
18
19use std::collections::HashMap;
20
21/// Milliseconds on the caller's monotonic clock.
22pub type Ms = u64;
23
24/// The most a source's interval is stretched for running over budget.
25const MAX_STRETCH: u32 = 4;
26/// Runs well inside the budget, in a row, before a stretch is halved.
27const CALM_RUNS: u32 = 3;
28
29/// The things that can be sampled.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
31pub enum SourceId {
32    /// CPU, memory, network and disk counters for the whole machine.
33    Machine,
34    /// Each pane's processes: CPU, memory and process count, and the cleanup of
35    /// panes whose shell is gone.
36    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/// What a source costs and needs.
49#[derive(Debug, Clone)]
50pub struct SourceSpec {
51    pub id: SourceId,
52    /// Never run more often than this, whatever a lease asks.
53    pub min_interval_ms: Ms,
54    /// How long one run should take. A longer one stretches the interval.
55    pub budget_ms: Ms,
56    /// Sources this one reads, which run first in the same pass. They must be
57    /// declared before it.
58    pub needs: &'static [SourceId],
59}
60
61#[derive(Debug, Clone)]
62struct Lease {
63    holder: String,
64    source: SourceId,
65    interval_ms: Ms,
66    /// `None`: held until released.
67    expires_at: Option<Ms>,
68}
69
70#[derive(Debug, Clone, Default)]
71struct SourceState {
72    /// When it runs next; `None` while nothing holds a lease on it.
73    next_due: Option<Ms>,
74    last_start: Option<Ms>,
75    last_took: Ms,
76    /// 1, 2 or 4: how many times its interval is stretched.
77    stretch: u32,
78    calm: u32,
79    runs: u64,
80    overruns: u64,
81    /// Runs that were due and didn't start for a whole interval or more.
82    skipped: u64,
83}
84
85/// What `sampler.status` reports for one source.
86#[derive(Debug, Clone, PartialEq, Eq)]
87pub struct SourceStatus {
88    pub id: SourceId,
89    pub holders: Vec<String>,
90    /// Between runs now, stretch included; `None` with no lease.
91    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    /// Sources due within this much of each other run in one pass, so a shared
104    /// input (one process snapshot) is read once for all of them.
105    coalesce_ms: Ms,
106}
107
108impl Scheduler {
109    /// `specs` in dependency order: each source's `needs` come before it.
110    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    /// `holder` wants `source` every `interval_ms` until `expires_at` (`None`:
123    /// until released). A holder has one lease per source, so asking again
124    /// renews it.
125    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    /// Drops all of `holder`'s leases.
137    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    /// Drops `holder`'s lease on `source`.
146    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    /// The leases changed, so the interval may have: a source that already ran
152    /// is next due one (new) interval after that run, sooner or later than it was
153    /// planned, and a pending run isn't kept for a viewer that has gone.
154    fn reschedule(&mut self, source: SourceId, now: Ms) {
155        let Some(interval) = self.effective_interval_ms(source, now) else {
156            // Nothing holds it now: forget the plan, so a lease that comes before the
157            // next `due` starts it at once instead of waiting out the old schedule.
158            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    /// Between runs now: the shortest live lease, floored at the source's
179    /// minimum, times its stretch. `None` with no live lease.
180    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    /// The sources to run now, dependencies first, each once. A source with no
187    /// live lease isn't run, and starts at once when it gets one.
188    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        // What a due source reads runs with it, leased or not.
205        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    /// When to call [`Self::due`] next; `None` when nothing is leased.
213    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    /// A run of `id` that started at `started` took `took_ms`. Schedules its next
222    /// run from its start, never sooner than its end, and never to make up for a
223    /// late one: a source that fell behind is counted as having skipped.
224    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    /// Runs whatever is due at `now`, each taking `took`, and returns what ran.
286    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        // A viewer arrives: the source is pulled in to the new interval.
320        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        // No extra run at the old fast schedule: its next is 30 s after its last.
327        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        // Released and taken again with no `due` in between, well before its run at 30 s.
347        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        // Machine is due at 1000, Panes at 1000 too; one arriving 30 ms late is
407        // still in the pass, one 200 ms late is not.
408        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        // Not leased itself, so it isn't in the status as wanted.
418        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); // budget 200
434            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        // Three quick runs in a row halve it, three more halve it again.
440        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); // over half the budget, under all of it
461            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        // The next was due at 1000 but the loop only got round to it at 4500.
472        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}