Skip to main content

grommet_core/
lib.rs

1//! Key-affine fair scheduling as a pure data structure.
2//!
3//! There is no async, no clock, no IO and no allocation policy here beyond the
4//! queues this module owns. Time is an explicit monotonic [`Duration`] supplied
5//! by the caller, and work items are opaque payloads. The exceptions are
6//! [`ring`] and [`waker_slot`], which are here for unsafe containment rather
7//! than because they are part of the scheduling model; see below.
8//!
9//! # Fairness model
10//!
11//! Each key has at most one in-flight item, so a key with a million queued
12//! items occupies exactly one dispatch slot, the same as a key with one. Given
13//! that, fairness reduces to *which ready key is dispatched next?*
14//!
15//! A completed key is never greedily re-dispatched. Both new arrivals and
16//! completions place a key at the BACK of a round-robin ready ring, and
17//! dispatch always pops from the FRONT. A key sitting at position `k` is
18//! therefore dispatched within `k` dispatches: a strict bound, hence
19//! starvation-free. One key cannot monopolize a shard.
20//!
21//! There is one ring per work class, each with its own in-flight budget, so a
22//! flood of one class cannot starve another. A key routes to a ring based on
23//! the class of its CURRENT queue head, so a key with mixed classes moves
24//! between rings in FIFO order as it drains.
25//!
26//! # Trusting the caller
27//!
28//! Everything this module knows about a payload: its affine key, its work
29//! class and its optional deadline: is stamped once by the caller at
30//! admission and never recomputed. A caller-supplied trait implementation that
31//! answered inconsistently on a second call would otherwise desynchronize the
32//! ready rings from the per-key queues, so the opportunity is removed rather
33//! than documented away.
34//!
35//! # Unsafe code
36//!
37//! This crate is where the workspace keeps its unsafe, so that there is one
38//! place to audit rather than a policy with exceptions. Every other crate is
39//! `#![deny(unsafe_code)]`, and here it is confined to modules that opt back in
40//! one at a time:
41//!
42//! - the **queue slab**, which indexes without bounds checks;
43//! - [`waker_slot`], which guards an `Option<Waker>` with a two-bit lock;
44//! - [`ring`], which guards a slot's value with a publication stamp;
45//! - the **cell** shim the last two share, which is one constructor and one
46//!   accessor and exists so that both present loom's checked `UnsafeCell` shape
47//!   rather than each rolling its own.
48//!
49//! Each states the invariant it relies on at the top of its file,
50//! `debug_assert!`s it at every use, and is checked by a model test under Miri.
51//! The two synchronization primitives are additionally model-checked by loom,
52//! which verifies the exclusion arguments mechanically: a slot read that
53//! escaped its stamp fails the model as a causality violation rather than
54//! passing as undefined behaviour that happens not to bite.
55
56#![deny(unsafe_code)]
57
58use ahash::AHashMap;
59#[cfg(not(coverage))]
60use grommet_macros::always;
61use std::collections::VecDeque;
62use std::collections::hash_map::Entry;
63use std::hash::Hash;
64use std::time::Duration;
65
66mod cell;
67
68pub mod timer;
69// Not part of the documented surface: this is a workspace-internal seam that
70// happens to need a crate boundary, and publishing it would commit this crate
71// to a synchronization primitive in its API.
72#[doc(hidden)]
73pub mod ring;
74#[doc(hidden)]
75pub mod waker_slot;
76
77mod queue;
78use queue::{List, Slab};
79
80/// Index of a work class, in `0..CLASSES`.
81pub type ClassId = u8;
82
83/// Per-shard scheduling limits.
84///
85/// `max_inflight` is indexed by [`ClassId`]: each class dispatches into its own
86/// budget, so saturating one class never blocks another.
87///
88/// Build one with [`Config::new`] and adjust the fields you care about. The
89/// struct is `#[non_exhaustive]` so that a later release can add a limit
90/// without that being a breaking change for anyone who did exactly that.
91#[derive(Clone, Copy, Debug, PartialEq, Eq)]
92#[non_exhaustive]
93pub struct Config<const CLASSES: usize = 2> {
94    /// Maximum simultaneously dispatched items per class.
95    pub max_inflight: [usize; CLASSES],
96    /// Maximum queued plus in-flight items. Admission above this is the
97    /// caller's responsibility to refuse; the value is exposed so the caller
98    /// can gate its own mailbox against it.
99    pub max_pending: usize,
100    /// Soft cap on resident keys. Eviction is bounded work per sweep, so the
101    /// cap can be exceeded transiently when every candidate is busy.
102    pub max_resident: Option<usize>,
103    /// How long a key must sit idle before it becomes an eviction candidate.
104    pub evict_after: Duration,
105    /// Maximum eviction candidates examined per sweep.
106    pub evict_iters: usize,
107    /// Queue slab entries reserved up front. The slab never exceeds the peak
108    /// number of simultaneously queued items, which `max_pending` bounds, so
109    /// reserving that much makes the steady state allocation-free at the cost
110    /// of the memory up front.
111    pub queue_reserve: usize,
112}
113
114impl<const CLASSES: usize> Config<CLASSES> {
115    /// A configuration with per-class in-flight budgets and defaults elsewhere.
116    pub fn new(max_inflight: [usize; CLASSES]) -> Self {
117        Self {
118            max_inflight,
119            max_pending: 8192,
120            max_resident: None,
121            evict_after: Duration::from_secs(60),
122            evict_iters: 256,
123            queue_reserve: 1024,
124        }
125    }
126}
127
128/// A work item being admitted, with its scheduling metadata already stamped.
129pub struct Admit<K, P> {
130    pub key: K,
131    pub class: ClassId,
132    /// Monotonic time after which this item is dropped instead of dispatched.
133    pub expires_at: Option<Duration>,
134    pub payload: P,
135}
136
137/// A work item that the scheduler has granted exclusive ownership of its key.
138pub struct Dispatch<K, P, S> {
139    pub key: K,
140    pub class: ClassId,
141    /// Resident state, moved out of the scheduler for the duration of the work.
142    /// `None` means the key has no state resident and the caller must load it.
143    pub state: Option<S>,
144    pub payload: P,
145}
146
147/// What became of a key's state once its work finished.
148#[derive(Clone, Copy, Debug, PartialEq, Eq)]
149pub enum Disposition<S> {
150    /// Keep this state resident for the next dispatch of the key.
151    Keep(S),
152    /// Discard any resident state: the next dispatch must reload it. This is
153    /// the correct answer whenever an operation's outcome is unknown, since a
154    /// stale in-memory value is worse than no value at all.
155    Drop,
156}
157
158impl<S> Disposition<S> {
159    pub fn into_option(self) -> Option<S> {
160        match self {
161            Self::Keep(state) => Some(state),
162            Self::Drop => None,
163        }
164    }
165}
166
167/// The result of a dispatched item, returned to the scheduler.
168pub struct Completion<K, S> {
169    pub key: K,
170    pub class: ClassId,
171    pub state: Disposition<S>,
172}
173
174/// Instantaneous scheduler gauges.
175#[derive(Clone, Copy, Debug, PartialEq, Eq)]
176pub struct Snapshot<const CLASSES: usize = 2> {
177    pub inflight: [usize; CLASSES],
178    pub ready: [usize; CLASSES],
179    pub pending: usize,
180    pub resident: usize,
181    pub evicting: usize,
182    /// Resident keys sitting idle, which is exactly the eviction sweep's
183    /// worklist. It is bounded by `resident` by construction; publishing it
184    /// is what makes that a claim the deployment can check rather than one
185    /// the source has to be trusted for.
186    pub eviction_backlog: usize,
187    /// High-water mark of simultaneously queued items, in slab entries. The
188    /// slab never shrinks, so this is what `Config::queue_reserve` should be
189    /// set to for an allocation-free steady state.
190    pub queue_capacity: usize,
191}
192
193// `[T; N]: Default` only covers `N <= 32`, so the class count cannot rely on it.
194impl<const CLASSES: usize> Default for Snapshot<CLASSES> {
195    fn default() -> Self {
196        Self {
197            inflight: [0; CLASSES],
198            ready: [0; CLASSES],
199            pending: 0,
200            resident: 0,
201            evicting: 0,
202            eviction_backlog: 0,
203            queue_capacity: 0,
204        }
205    }
206}
207
208#[derive(Clone, Copy, PartialEq, Eq, Debug)]
209enum Presence {
210    Idle,
211    Ready(ClassId),
212    InFlight,
213    /// State has been handed to the caller for flushing. The key is quiesced:
214    /// arriving work queues behind the flush rather than racing it.
215    Evicting,
216}
217
218struct Item<P> {
219    class: ClassId,
220    expires_at: Option<Duration>,
221    payload: P,
222}
223
224/// Per-key bookkeeping. Its size is independent of the work item type, because
225/// queued items live in the shard-wide slab rather than in the slot, which
226/// keeps the key map compact when a shard holds many keys.
227struct Slot<K, S> {
228    resident: Option<S>,
229    queue: List,
230    presence: Presence,
231    /// When this key last went idle, which is when its eviction window opens.
232    /// Read only while `presence` is [`Presence::Idle`]; stale otherwise, and
233    /// deliberately not maintained there, because nothing may consult it.
234    idle_since: Duration,
235    /// Neighbours in the shard's idle list, oldest first. Meaningful only
236    /// while `presence` is [`Presence::Idle`], and `None` in every other
237    /// state: a key that is ready, running or flushing is not a candidate.
238    idle_prev: Option<K>,
239    idle_next: Option<K>,
240}
241
242impl<K, S> Slot<K, S> {
243    /// A slot for a key seen for the first time. It is deliberately *not* on
244    /// the idle list: the caller creates one only to queue work behind it, so
245    /// it leaves `Idle` before the next statement observes it.
246    fn cold() -> Self {
247        Self {
248            resident: None,
249            queue: List::default(),
250            presence: Presence::Idle,
251            idle_since: Duration::ZERO,
252            idle_prev: None,
253            idle_next: None,
254        }
255    }
256
257    /// Take this slot's links, leaving it detached. The caller still owes the
258    /// neighbours their patch, which is what [`Scheduler::idle_patch`] does.
259    fn detach(&mut self) -> (Option<K>, Option<K>) {
260        (self.idle_prev.take(), self.idle_next.take())
261    }
262}
263
264/// A key-affine, class-fair scheduler over opaque payloads `P` and per-key
265/// state `S`.
266pub struct Scheduler<K, P, S, const CLASSES: usize = 2> {
267    cfg: Config<CLASSES>,
268    keys: AHashMap<K, Slot<K, S>>,
269    slab: Slab<Item<P>>,
270    ready: [VecDeque<K>; CLASSES],
271    /// Oldest and newest key on the idle list, an intrusive doubly-linked list
272    /// threaded through the slots themselves.
273    ///
274    /// Keys join at the tail when they go idle and are unlinked the moment
275    /// they are touched, so the list holds each idle key exactly once and
276    /// nothing else: it is bounded by the resident key count, and a sweep
277    /// walking from the head sees keys in genuine least-recently-used order
278    /// with no stale entries to skip.
279    idle_head: Option<K>,
280    idle_tail: Option<K>,
281    idle: usize,
282    expired: VecDeque<(K, ClassId, P)>,
283    inflight: [usize; CLASSES],
284    pending: usize,
285    evicting: usize,
286}
287
288impl<K, P, S, const CLASSES: usize> Scheduler<K, P, S, CLASSES>
289where
290    K: Copy + Eq + Hash,
291{
292    pub fn new(cfg: Config<CLASSES>) -> Self {
293        Self {
294            cfg,
295            keys: AHashMap::new(),
296            slab: Slab::with_capacity(cfg.queue_reserve),
297            ready: std::array::from_fn(|_| VecDeque::new()),
298            idle_head: None,
299            idle_tail: None,
300            idle: 0,
301            expired: VecDeque::new(),
302            inflight: [0; CLASSES],
303            pending: 0,
304            evicting: 0,
305        }
306    }
307
308    pub fn config(&self) -> &Config<CLASSES> {
309        &self.cfg
310    }
311
312    pub fn max_pending(&self) -> usize {
313        self.cfg.max_pending
314    }
315
316    /// Queued plus in-flight items across every class.
317    pub fn pending(&self) -> usize {
318        self.pending
319    }
320
321    pub fn is_saturated(&self) -> bool {
322        self.pending >= self.cfg.max_pending
323    }
324
325    /// Queue an item behind its key. The caller is responsible for refusing
326    /// admission above [`Config::max_pending`]; this is where backpressure is
327    /// applied, and the scheduler deliberately does not decide the policy.
328    ///
329    /// Admission takes no clock. A key's eviction window opens when it goes
330    /// idle, not when work arrives for it, so there is no time to record here.
331    pub fn admit(&mut self, item: Admit<K, P>) {
332        let Admit { key, class, expires_at, payload } = item;
333        debug_assert!((class as usize) < CLASSES, "class {class} is outside 0..{CLASSES}");
334        self.pending += 1;
335        // Whether the slot already existed decides whether it is on the idle
336        // list, and the entry answers that without a second lookup.
337        let (slot, resurrected) = match self.keys.entry(key) {
338            Entry::Occupied(entry) => (entry.into_mut(), true),
339            Entry::Vacant(entry) => (entry.insert(Slot::cold()), false),
340        };
341        self.slab.push_back(&mut slot.queue, Item { class, expires_at, payload });
342        // A key that is in-flight, already ready, or quiescing for eviction
343        // keeps its position; only an idle key joins a ring, and its queue was
344        // empty, so the item just pushed is the head whose class decides.
345        let mut detached = None;
346        let joins = match slot.presence {
347            Presence::Idle => {
348                slot.presence = Presence::Ready(class);
349                // An idle key that was already resident is on the idle list
350                // and has just stopped being a candidate. One created a
351                // statement ago never joined it.
352                if resurrected {
353                    detached = Some(slot.detach());
354                }
355                true
356            }
357            Presence::Ready(_) | Presence::InFlight | Presence::Evicting => false,
358        };
359        if let Some((prev, next)) = detached {
360            self.idle -= 1;
361            self.idle_patch(prev, next);
362        }
363        if joins {
364            self.ready[class as usize].push_back(key);
365        }
366    }
367
368    /// Dispatch the next ready item of `class`, taking exclusive ownership of
369    /// its key. Items whose deadline has passed are discarded on the way and
370    /// can be collected with [`Scheduler::pop_expired`].
371    pub fn next(&mut self, class: ClassId, now: Duration) -> Option<Dispatch<K, P, S>> {
372        let index = class as usize;
373        if self.inflight[index] >= self.cfg.max_inflight[index] {
374            return None;
375        }
376        loop {
377            let key = self.ready[index].pop_front()?;
378            let slot = self.keys.get_mut(&key).expect("ready key has a slot");
379
380            // Drop expired heads. Disjoint field borrows keep this legal while
381            // `slot` is live, so no scratch buffer is needed.
382            let mut taken = None;
383            while let Some(head) = self.slab.front(&slot.queue) {
384                if head.class != class {
385                    break;
386                }
387                if head.expires_at.is_some_and(|deadline| deadline <= now) {
388                    let item =
389                        self.slab.pop_front(&mut slot.queue).expect("front was just observed");
390                    self.pending -= 1;
391                    self.expired.push_back((key, item.class, item.payload));
392                    continue;
393                }
394                taken = self.slab.pop_front(&mut slot.queue);
395                break;
396            }
397
398            match taken {
399                Some(item) => {
400                    slot.presence = Presence::InFlight;
401                    let state = slot.resident.take();
402                    self.inflight[index] += 1;
403                    return Some(Dispatch { key, class, state, payload: item.payload });
404                }
405                None => {
406                    // Everything of this class expired. The key now heads a
407                    // different ring, or has nothing left at all. Either way it
408                    // cannot rejoin the ring being drained, so this terminates.
409                    let target = Self::settle(&self.slab, slot);
410                    debug_assert!(target != Some(class));
411                    self.place(key, target, now);
412                }
413            }
414        }
415    }
416
417    /// Take an item that was discarded at dispatch because its deadline had
418    /// passed, with the class it was admitted under. The caller owns telling
419    /// whoever submitted it.
420    ///
421    /// The class comes from the item rather than being recomputed from the
422    /// payload, for the reason the module header gives: what this scheduler
423    /// knows about a payload is stamped once at admission, and asking again
424    /// would let an inconsistent answer desynchronize the accounting.
425    pub fn pop_expired(&mut self) -> Option<(K, ClassId, P)> {
426        self.expired.pop_front()
427    }
428
429    /// Return a key's state and re-place it in its ring, or retire it.
430    pub fn complete(&mut self, completion: Completion<K, S>, now: Duration) {
431        let Completion { key, class, state } = completion;
432        self.inflight[class as usize] -= 1;
433        self.pending -= 1;
434        let slot = self.keys.get_mut(&key).expect("completed key has a slot");
435        debug_assert_eq!(slot.presence, Presence::InFlight);
436        slot.resident = state.into_option();
437        let target = Self::settle(&self.slab, slot);
438        self.place(key, target, now);
439    }
440
441    /// Collect keys whose resident state should be flushed and released,
442    /// appending `(key, state)` pairs to `out`. Each collected key is quiesced
443    /// until [`Scheduler::finish_evict`] is called for it: work may queue
444    /// behind the flush, but nothing dispatches, so a write-back flush cannot
445    /// race a reload of the same key.
446    pub fn evict(&mut self, now: Duration, out: &mut Vec<(K, S)>) {
447        for _ in 0..self.cfg.evict_iters {
448            // The head is the least recently idled key, so the first one that
449            // is too young ends the sweep: nothing behind it can be older.
450            let Some(key) = self.idle_head else {
451                break;
452            };
453            let since = self.keys.get(&key).expect("a listed key has a slot").idle_since;
454            let idle_long_enough = now.saturating_sub(since) >= self.cfg.evict_after;
455            // A quiescing key still occupies a map entry but its state is
456            // already being flushed, so it must not count against the cap or a
457            // single sweep would evict far past it.
458            let resident = self.keys.len() - self.evicting;
459            let over_capacity = self.cfg.max_resident.is_some_and(|max| resident > max);
460            if !idle_long_enough && !over_capacity {
461                break;
462            }
463            self.idle_unlink_head();
464            self.release(key, out);
465        }
466    }
467
468    /// Flush and release every resident key at once, whatever its idle time and
469    /// whatever the capacity cap says, ignoring [`Config::evict_iters`].
470    ///
471    /// This is the shutdown path. Resident state is a write-back cache: the
472    /// scheduler holds the only copy between dispatches, so a shard that stops
473    /// without draining it discards writes the processor was told it could keep.
474    /// The keys are quiesced exactly as [`Scheduler::evict`] quiesces them, so
475    /// each one still needs its [`Scheduler::finish_evict`].
476    pub fn evict_all(&mut self, out: &mut Vec<(K, S)>) {
477        while let Some(key) = self.idle_head {
478            self.idle_unlink_head();
479            self.release(key, out);
480        }
481    }
482
483    /// Quiesce one eviction candidate and hand back its state to flush, or drop
484    /// the key outright when it has none.
485    ///
486    /// The key must already be off the idle list, which is what makes this
487    /// unconditional: the list holds only idle keys, exactly once each, so
488    /// there is no stale candidate left to recognize and skip.
489    fn release(&mut self, key: K, out: &mut Vec<(K, S)>) {
490        let slot = self.keys.get_mut(&key).expect("a listed key has a slot");
491        debug_assert_eq!(slot.presence, Presence::Idle);
492        match slot.resident.take() {
493            Some(state) => {
494                slot.presence = Presence::Evicting;
495                self.evicting += 1;
496                out.push((key, state));
497            }
498            None => {
499                self.keys.remove(&key);
500            }
501        }
502    }
503
504    /// Release a key quiesced by [`Scheduler::evict`] once its flush finished.
505    pub fn finish_evict(&mut self, key: K, now: Duration) {
506        self.evicting -= 1;
507        let Some(slot) = self.keys.get_mut(&key) else {
508            debug_assert!(false, "finished eviction for an unknown key");
509            return;
510        };
511        debug_assert_eq!(slot.presence, Presence::Evicting);
512        if slot.queue.is_empty() {
513            self.keys.remove(&key);
514            return;
515        }
516        // Work arrived during the flush. The state is gone, so the next
517        // dispatch reloads it.
518        let target = Self::settle(&self.slab, slot);
519        self.place(key, target, now);
520    }
521
522    pub fn snapshot(&self) -> Snapshot<CLASSES> {
523        Snapshot {
524            inflight: self.inflight,
525            ready: std::array::from_fn(|class| self.ready[class].len()),
526            pending: self.pending,
527            resident: self.keys.len(),
528            evicting: self.evicting,
529            eviction_backlog: self.idle,
530            queue_capacity: self.slab.capacity(),
531        }
532    }
533
534    /// Set a slot's presence from its queue head, returning the ring it should
535    /// join, or `None` when it has become idle.
536    fn settle(slab: &Slab<Item<P>>, slot: &mut Slot<K, S>) -> Option<ClassId> {
537        match slab.front(&slot.queue) {
538            Some(head) => {
539                slot.presence = Presence::Ready(head.class);
540                Some(head.class)
541            }
542            None => {
543                slot.presence = Presence::Idle;
544                None
545            }
546        }
547    }
548
549    fn place(&mut self, key: K, target: Option<ClassId>, now: Duration) {
550        match target {
551            Some(class) => self.ready[class as usize].push_back(key),
552            None => self.idle_link(key, now),
553        }
554    }
555
556    /// Append a key that has just gone idle to the newest end of the idle
557    /// list, stamping the moment its eviction window opened.
558    fn idle_link(&mut self, key: K, now: Duration) {
559        let tail = self.idle_tail;
560        let slot = self.keys.get_mut(&key).expect("a settled key has a slot");
561        debug_assert_eq!(slot.presence, Presence::Idle);
562        debug_assert!(slot.idle_prev.is_none() && slot.idle_next.is_none());
563        slot.idle_since = now;
564        slot.idle_prev = tail;
565        match tail {
566            Some(previous) => {
567                self.keys.get_mut(&previous).expect("a listed key has a slot").idle_next =
568                    Some(key);
569            }
570            None => self.idle_head = Some(key),
571        }
572        self.idle_tail = Some(key);
573        self.idle += 1;
574    }
575
576    /// Detach the oldest idle key. The caller must have observed a non-empty
577    /// list, and owns whatever becomes of the key afterwards.
578    fn idle_unlink_head(&mut self) {
579        let key = self.idle_head.expect("the caller observed a listed key");
580        let (prev, next) = self.keys.get_mut(&key).expect("a listed key has a slot").detach();
581        debug_assert!(prev.is_none(), "the head of the idle list has no predecessor");
582        self.idle -= 1;
583        self.idle_patch(prev, next);
584    }
585
586    /// Close the idle list over a slot that was just detached from it.
587    fn idle_patch(&mut self, prev: Option<K>, next: Option<K>) {
588        match prev {
589            Some(key) => {
590                self.keys.get_mut(&key).expect("a listed key has a slot").idle_next = next;
591            }
592            None => self.idle_head = next,
593        }
594        match next {
595            Some(key) => {
596                self.keys.get_mut(&key).expect("a listed key has a slot").idle_prev = prev;
597            }
598            None => self.idle_tail = prev,
599        }
600    }
601
602    #[cfg(coverage)]
603    pub fn check_invariants(&self) -> Result<(), &'static str> {
604        Ok(())
605    }
606
607    /// Verify every structural invariant. Linear in resident keys plus total
608    /// ring length, with no allocation, so it is affordable under a debug
609    /// assertion on every reactor turn, which is where it earns its keep.
610    #[cfg(not(coverage))]
611    pub fn check_invariants(&self) -> Result<(), &'static str> {
612        let queued: usize = self.keys.values().map(|slot| slot.queue.len()).sum();
613        let inflight: usize = self.inflight.iter().sum();
614        if !always!(self.pending == queued + inflight) {
615            return Err("pending != queued + in-flight");
616        }
617        let owned = self.keys.values().filter(|slot| slot.presence == Presence::InFlight).count();
618        if !always!(owned == inflight) {
619            return Err("in-flight counters disagree with key ownership");
620        }
621        let quiesced =
622            self.keys.values().filter(|slot| slot.presence == Presence::Evicting).count();
623        if !always!(quiesced == self.evicting) {
624            return Err("evicting counter disagrees with key ownership");
625        }
626        // Ring membership is checked outward from the rings and then reconciled
627        // by count, rather than by asking each key which rings hold it. The
628        // latter reads every ring once per key: quadratic, and at a few
629        // thousand resident keys it costs more per item than the scheduling it
630        // is guarding, which put a ceiling on how large a simulation could run.
631        let mut listed = 0;
632        for (class, ring) in self.ready.iter().enumerate() {
633            let class = class as ClassId;
634            for key in ring {
635                listed += 1;
636                let Some(slot) = self.keys.get(key) else {
637                    return Err("a ready key has no slot");
638                };
639                if slot.presence != Presence::Ready(class) {
640                    return Err("ready key is inconsistent with the ring holding it");
641                }
642                // Covers an empty queue too: it has no head to disagree.
643                if self.slab.front(&slot.queue).map(|item| item.class) != Some(class) {
644                    return Err("ready key is inconsistent with its queue head");
645                }
646            }
647        }
648        // A key listed in two rings, or twice in one, would need a second
649        // `Ready` slot to balance this, and a slot has one presence.
650        let ready =
651            self.keys.values().filter(|slot| matches!(slot.presence, Presence::Ready(_))).count();
652        if !always!(listed == ready) {
653            return Err("ready rings disagree with key presence");
654        }
655
656        for slot in self.keys.values() {
657            match slot.presence {
658                Presence::Idle if !slot.queue.is_empty() => {
659                    return Err("idle key is queued");
660                }
661                Presence::InFlight | Presence::Evicting if slot.resident.is_some() => {
662                    return Err("a key that gave up its state still holds it");
663                }
664                _ => {}
665            }
666        }
667
668        // The idle list is what bounds this scheduler's memory: it must hold
669        // every idle key, hold nothing else, and hold each of them once. The
670        // walk below establishes all three: forward links reach only idle
671        // keys, back links agree with them (so no key appears twice without
672        // one of the two disagreeing), and the count reconciles against the
673        // number of idle slots, which catches a key that is simply missing.
674        let mut walked = 0;
675        let mut previous = None;
676        let mut cursor = self.idle_head;
677        while let Some(key) = cursor {
678            // A cycle would otherwise spin here forever rather than fail.
679            if walked > self.keys.len() {
680                return Err("the idle list cycles");
681            }
682            walked += 1;
683            let Some(slot) = self.keys.get(&key) else {
684                return Err("a key on the idle list has no slot");
685            };
686            if slot.presence != Presence::Idle {
687                return Err("a key that is not idle is on the idle list");
688            }
689            if slot.idle_prev != previous {
690                return Err("the idle list's back links disagree with its forward links");
691            }
692            previous = cursor;
693            cursor = slot.idle_next;
694        }
695        if self.idle_tail != previous {
696            return Err("the idle list's tail is not the last key on it");
697        }
698        if !always!(walked == self.idle) {
699            return Err("the idle counter disagrees with the idle list");
700        }
701        let idle = self.keys.values().filter(|slot| slot.presence == Presence::Idle).count();
702        if !always!(walked == idle) {
703            return Err("an idle key is missing from the idle list");
704        }
705        Ok(())
706    }
707}
708
709#[cfg(test)]
710mod tests {
711    use super::*;
712
713    const IO: ClassId = 0;
714    const CPU: ClassId = 1;
715
716    type Book = Scheduler<u64, &'static str, u64, 2>;
717
718    fn config() -> Config<2> {
719        Config {
720            max_inflight: [1, 1],
721            max_pending: 32,
722            max_resident: None,
723            evict_after: Duration::from_secs(10),
724            evict_iters: 32,
725            queue_reserve: 8,
726        }
727    }
728
729    fn item(key: u64, class: ClassId) -> Admit<u64, &'static str> {
730        Admit { key, class, expires_at: None, payload: "work" }
731    }
732
733    fn expiring(key: u64, class: ClassId, at: Duration) -> Admit<u64, &'static str> {
734        Admit { key, class, expires_at: Some(at), payload: "work" }
735    }
736
737    fn finish(dispatch: Dispatch<u64, &'static str, u64>) -> Completion<u64, u64> {
738        Completion {
739            key: dispatch.key,
740            class: dispatch.class,
741            state: Disposition::Keep(dispatch.state.unwrap_or_default()),
742        }
743    }
744
745    #[test]
746    fn a_backlogged_key_rotates_behind_every_other_ready_key() {
747        let mut book = Book::new(config());
748        let now = Duration::ZERO;
749        book.admit(item(1, IO));
750        book.admit(item(1, IO));
751        book.admit(item(2, IO));
752
753        let whale = book.next(IO, now).unwrap();
754        assert_eq!(whale.key, 1);
755        book.complete(finish(whale), now);
756        assert_eq!(book.next(IO, now).unwrap().key, 2, "the backlog must not be served twice");
757        assert_eq!(book.check_invariants(), Ok(()));
758    }
759
760    #[test]
761    fn dispatch_position_bounds_starvation_under_sustained_load() {
762        let mut book = Book::new(config());
763        let now = Duration::ZERO;
764        for key in 0..8 {
765            book.admit(item(key, IO));
766        }
767        // The hot key keeps arriving; the strict bound says every other key is
768        // still served within one rotation of the ring.
769        let mut seen = [false; 8];
770        for _ in 0..8 {
771            book.admit(item(0, IO));
772            let dispatch = book.next(IO, now).unwrap();
773            seen[dispatch.key as usize] = true;
774            book.complete(finish(dispatch), now);
775        }
776        assert!(seen.into_iter().all(|served| served), "a key starved behind the hot key");
777    }
778
779    #[test]
780    fn class_budgets_are_independent_and_a_key_serializes_across_them() {
781        let mut book = Book::new(config());
782        let now = Duration::ZERO;
783        book.admit(item(1, IO));
784        book.admit(item(1, IO));
785        book.admit(item(2, CPU));
786
787        let io = book.next(IO, now).unwrap();
788        let cpu = book.next(CPU, now).unwrap();
789        assert!(book.next(IO, now).is_none(), "key 1 already owns its single in-flight slot");
790        assert!(book.next(CPU, now).is_none(), "the compute budget is saturated");
791        assert_eq!(book.check_invariants(), Ok(()));
792
793        book.complete(finish(io), now);
794        assert!(book.next(IO, now).is_some());
795        assert!(book.next(CPU, now).is_none(), "completing IO must not free compute budget");
796        book.complete(finish(cpu), now);
797    }
798
799    #[test]
800    fn a_mixed_key_moves_between_rings_in_fifo_order() {
801        let mut book = Book::new(config());
802        let now = Duration::ZERO;
803        book.admit(item(9, IO));
804        book.admit(item(9, CPU));
805        assert!(book.next(CPU, now).is_none(), "the compute item is behind the IO item");
806
807        let first = book.next(IO, now).unwrap();
808        book.complete(finish(first), now);
809        assert_eq!(book.next(CPU, now).unwrap().key, 9);
810    }
811
812    #[test]
813    fn state_ownership_transfers_to_exactly_one_dispatch() {
814        let mut book = Book::new(config());
815        let now = Duration::ZERO;
816        book.admit(item(3, IO));
817        let first = book.next(IO, now).unwrap();
818        assert_eq!(first.state, None, "a cold key carries no state");
819        book.complete(Completion { key: 3, class: IO, state: Disposition::Keep(77) }, now);
820
821        book.admit(item(3, IO));
822        let second = book.next(IO, now).unwrap();
823        assert_eq!(second.state, Some(77), "resident state follows the key");
824        book.complete(Completion { key: 3, class: IO, state: Disposition::Drop }, now);
825
826        book.admit(item(3, IO));
827        let third = book.next(IO, now).unwrap();
828        assert_eq!(third.state, None, "a dropped disposition forces a reload");
829        book.complete(finish(third), now);
830    }
831
832    #[test]
833    fn expired_items_are_discarded_at_dispatch_and_handed_back() {
834        let mut book = Book::new(config());
835        let deadline = Duration::from_secs(1);
836        book.admit(expiring(4, IO, deadline));
837        book.admit(expiring(5, IO, deadline));
838        book.admit(item(6, IO));
839
840        let now = Duration::from_secs(2);
841        let dispatch = book.next(IO, now).expect("the item without a deadline survives");
842        assert_eq!(dispatch.key, 6);
843        assert_eq!(book.pop_expired().map(|(key, ..)| key), Some(4));
844        assert_eq!(book.pop_expired().map(|(key, ..)| key), Some(5));
845        assert_eq!(book.pop_expired().map(|(key, ..)| key), None);
846        assert_eq!(book.pending(), 1, "expired items leave the pending count");
847        book.complete(finish(dispatch), now);
848        assert_eq!(book.check_invariants(), Ok(()));
849    }
850
851    #[test]
852    fn expiring_a_ring_head_re_places_the_key_on_its_next_class() {
853        let mut book = Book::new(config());
854        let deadline = Duration::from_secs(1);
855        book.admit(expiring(7, IO, deadline));
856        book.admit(item(7, CPU));
857
858        let now = Duration::from_secs(2);
859        assert!(book.next(IO, now).is_none(), "the only IO item expired");
860        assert_eq!(book.pop_expired().map(|(key, ..)| key), Some(7));
861        assert_eq!(book.next(CPU, now).unwrap().key, 7, "the key moved to the compute ring");
862        assert_eq!(book.check_invariants(), Ok(()));
863    }
864
865    #[test]
866    fn idle_keys_are_evicted_after_their_ttl_and_flushed_once() {
867        let mut book = Book::new(config());
868        book.admit(item(5, IO));
869        let _dispatch = book.next(IO, Duration::ZERO).unwrap();
870        book.complete(
871            Completion { key: 5, class: IO, state: Disposition::Keep(42) },
872            Duration::ZERO,
873        );
874
875        let mut flushed = Vec::new();
876        book.evict(Duration::from_secs(5), &mut flushed);
877        assert!(flushed.is_empty(), "the key is still inside its idle window");
878
879        book.evict(Duration::from_secs(11), &mut flushed);
880        assert_eq!(flushed, vec![(5, 42)]);
881        assert_eq!(book.snapshot().evicting, 1);
882        assert_eq!(book.check_invariants(), Ok(()));
883
884        book.finish_evict(5, Duration::from_secs(11));
885        assert_eq!(book.snapshot().resident, 0);
886    }
887
888    #[test]
889    fn work_arriving_during_a_flush_waits_and_then_reloads() {
890        let mut book = Book::new(config());
891        book.admit(item(8, IO));
892        let _dispatch = book.next(IO, Duration::ZERO).unwrap();
893        book.complete(
894            Completion { key: 8, class: IO, state: Disposition::Keep(11) },
895            Duration::ZERO,
896        );
897
898        let mut flushed = Vec::new();
899        let now = Duration::from_secs(11);
900        book.evict(now, &mut flushed);
901        assert_eq!(flushed, vec![(8, 11)]);
902
903        // The key is quiesced: nothing dispatches while the flush is running.
904        book.admit(item(8, IO));
905        assert!(book.next(IO, now).is_none(), "a quiesced key must not dispatch");
906        assert_eq!(book.check_invariants(), Ok(()));
907
908        book.finish_evict(8, now);
909        let after = book.next(IO, now).expect("the key resumes once the flush completes");
910        assert_eq!(after.state, None, "flushed state is never silently reused");
911        book.complete(finish(after), now);
912    }
913
914    #[test]
915    fn evict_all_flushes_every_resident_key_regardless_of_idle_time() {
916        let mut book = Book::new(Config { evict_iters: 1, ..config() });
917        for key in 0..4 {
918            book.admit(item(key, IO));
919            let _dispatch = book.next(IO, Duration::ZERO).expect("the key dispatches");
920            book.complete(
921                Completion { key, class: IO, state: Disposition::Keep(key * 10) },
922                Duration::ZERO,
923            );
924        }
925
926        let mut flushed = Vec::new();
927        // Well inside the idle window, and `evict_iters` would cap a sweep at
928        // one key, so neither limit may apply on the shutdown path.
929        book.evict_all(&mut flushed);
930        assert_eq!(flushed, vec![(0, 0), (1, 10), (2, 20), (3, 30)]);
931        assert_eq!(book.snapshot().evicting, 4);
932        assert_eq!(book.check_invariants(), Ok(()));
933
934        for (key, _) in flushed {
935            book.finish_evict(key, Duration::ZERO);
936        }
937        assert_eq!(book.snapshot().resident, 0, "every key is released once its flush lands");
938        assert_eq!(book.check_invariants(), Ok(()));
939    }
940
941    #[test]
942    fn evict_all_skips_keys_that_are_not_idle() {
943        let mut book = Book::new(config());
944        // Key 2 goes resident and idle; key 1 is left holding its state.
945        book.admit(item(2, IO));
946        let _dispatch = book.next(IO, Duration::ZERO).expect("key 2 dispatches");
947        book.complete(
948            Completion { key: 2, class: IO, state: Disposition::Keep(9) },
949            Duration::ZERO,
950        );
951        book.admit(item(1, IO));
952        let inflight = book.next(IO, Duration::ZERO).expect("key 1 dispatches");
953
954        let mut flushed = Vec::new();
955        book.evict_all(&mut flushed);
956        assert_eq!(flushed, vec![(2, 9)], "an in-flight key still owns its state");
957        assert_eq!(book.check_invariants(), Ok(()));
958        book.complete(finish(inflight), Duration::ZERO);
959    }
960
961    #[test]
962    fn reactivating_a_key_restarts_its_idle_window_from_the_back_of_the_list() {
963        let mut book = Book::new(config());
964        book.admit(item(5, IO));
965        let first = book.next(IO, Duration::ZERO).unwrap();
966        book.complete(finish(first), Duration::ZERO);
967        assert_eq!(book.snapshot().eviction_backlog, 1);
968
969        let later = Duration::from_secs(5);
970        book.admit(item(5, IO));
971        assert_eq!(book.snapshot().eviction_backlog, 0, "a touched key leaves the idle list");
972        let second = book.next(IO, later).unwrap();
973        book.complete(finish(second), later);
974        assert_eq!(book.snapshot().eviction_backlog, 1, "and rejoins it when it settles");
975        assert_eq!(book.check_invariants(), Ok(()));
976
977        let mut flushed = Vec::new();
978        book.evict(Duration::from_secs(11), &mut flushed);
979        assert!(flushed.is_empty(), "the window runs from the second completion, not the first");
980        assert_eq!(book.snapshot().resident, 1);
981
982        book.evict(Duration::from_secs(16), &mut flushed);
983        assert_eq!(flushed.len(), 1);
984    }
985
986    /// The property the intrusive list exists for. A key that cycles through
987    /// idle over and over must occupy exactly one entry the whole time: under
988    /// the timestamped candidate queue this grew once per completion, without
989    /// any bound, which is the bug this replaced.
990    #[test]
991    fn idle_tracking_is_bounded_by_resident_keys_however_often_they_cycle() {
992        let mut book = Book::new(Config { evict_iters: 0, ..config() });
993        let mut now = Duration::ZERO;
994        for round in 0..500u64 {
995            for key in 0..4 {
996                now += Duration::from_millis(1);
997                book.admit(item(key, IO));
998                let dispatch = book.next(IO, now).expect("the key dispatches");
999                book.complete(finish(dispatch), now);
1000            }
1001            let snapshot = book.snapshot();
1002            assert_eq!(
1003                snapshot.eviction_backlog, snapshot.resident,
1004                "every resident key is idle here, and each may appear once (round {round})"
1005            );
1006            assert_eq!(snapshot.eviction_backlog, 4, "four keys, whatever the throughput");
1007        }
1008        assert_eq!(book.check_invariants(), Ok(()));
1009    }
1010
1011    /// Eviction order is by idle time, not by arrival, so touching the middle
1012    /// of the list has to move that key to the back of it.
1013    #[test]
1014    fn the_idle_list_evicts_least_recently_idled_first_across_unlink_positions() {
1015        let mut book = Book::new(Config { max_resident: Some(0), evict_iters: 8, ..config() });
1016        fn settle(book: &mut Book, key: u64, now: Duration) {
1017            book.admit(item(key, IO));
1018            book.next(IO, now).expect("the key dispatches");
1019            book.complete(Completion { key, class: IO, state: Disposition::Keep(key) }, now);
1020        }
1021        for key in 0..4 {
1022            settle(&mut book, key, Duration::from_secs(key));
1023        }
1024        // Touch the head and the middle, in that order: both unlink from
1025        // positions that exercise different arms of the patch.
1026        settle(&mut book, 0, Duration::from_secs(10));
1027        settle(&mut book, 2, Duration::from_secs(11));
1028        assert_eq!(book.check_invariants(), Ok(()));
1029
1030        let mut flushed = Vec::new();
1031        book.evict(Duration::from_secs(12), &mut flushed);
1032        assert_eq!(
1033            flushed,
1034            vec![(1, 1), (3, 3), (0, 0), (2, 2)],
1035            "untouched keys first in their original order, then the two that were touched"
1036        );
1037        assert_eq!(book.snapshot().eviction_backlog, 0);
1038    }
1039
1040    #[test]
1041    fn a_key_with_no_resident_state_leaves_the_idle_list_by_being_dropped() {
1042        let mut book = Book::new(config());
1043        book.admit(item(6, IO));
1044        let dispatch = book.next(IO, Duration::ZERO).unwrap();
1045        book.complete(Completion { key: 6, class: IO, state: Disposition::Drop }, Duration::ZERO);
1046        let _ = dispatch;
1047        assert_eq!(
1048            book.snapshot().eviction_backlog,
1049            1,
1050            "a stateless idle key is still a candidate"
1051        );
1052
1053        let mut flushed = Vec::new();
1054        book.evict(Duration::from_secs(11), &mut flushed);
1055        assert!(flushed.is_empty(), "there is nothing to flush");
1056        let snapshot = book.snapshot();
1057        assert_eq!(snapshot.resident, 0, "but the map entry is reclaimed");
1058        assert_eq!(snapshot.eviction_backlog, 0);
1059        assert_eq!(book.check_invariants(), Ok(()));
1060    }
1061
1062    #[test]
1063    fn a_key_quiescing_for_eviction_is_not_a_candidate_again_until_it_settles() {
1064        let mut book = Book::new(config());
1065        book.admit(item(7, IO));
1066        let dispatch = book.next(IO, Duration::ZERO).unwrap();
1067        book.complete(
1068            Completion { key: 7, class: IO, state: Disposition::Keep(3) },
1069            Duration::ZERO,
1070        );
1071        let _ = dispatch;
1072
1073        let mut flushed = Vec::new();
1074        let now = Duration::from_secs(11);
1075        book.evict(now, &mut flushed);
1076        assert_eq!(flushed, vec![(7, 3)]);
1077        assert_eq!(book.snapshot().eviction_backlog, 0, "an evicting key is off the list");
1078
1079        // Work queues behind the flush, so the key settles into a ring rather
1080        // than back onto the idle list.
1081        book.admit(item(7, IO));
1082        book.finish_evict(7, now);
1083        assert_eq!(book.snapshot().eviction_backlog, 0);
1084        let after = book.next(IO, now).expect("the key resumes");
1085        book.complete(finish(after), now);
1086        assert_eq!(book.snapshot().eviction_backlog, 1, "and only rejoins once it is idle");
1087        assert_eq!(book.check_invariants(), Ok(()));
1088    }
1089
1090    #[test]
1091    fn capacity_pressure_evicts_before_the_idle_window_elapses() {
1092        let mut book = Book::new(Config { max_resident: Some(1), ..config() });
1093        for key in 0..3 {
1094            book.admit(item(key, IO));
1095            let _dispatch = book.next(IO, Duration::ZERO).unwrap();
1096            book.complete(
1097                Completion { key, class: IO, state: Disposition::Keep(key) },
1098                Duration::ZERO,
1099            );
1100        }
1101        assert_eq!(book.snapshot().resident, 3);
1102
1103        let mut flushed = Vec::new();
1104        book.evict(Duration::ZERO, &mut flushed);
1105        assert_eq!(flushed, vec![(0, 0), (1, 1)], "the oldest idle keys go first");
1106        for (key, _) in flushed {
1107            book.finish_evict(key, Duration::ZERO);
1108        }
1109        assert_eq!(book.snapshot().resident, 1);
1110        assert_eq!(book.check_invariants(), Ok(()));
1111    }
1112
1113    #[test]
1114    fn snapshot_and_saturation_report_the_scheduler_state() {
1115        let mut book = Book::new(Config { max_pending: 2, ..config() });
1116        assert_eq!(book.snapshot(), Snapshot::default());
1117        assert_eq!(book.max_pending(), 2);
1118        assert_eq!(book.config().max_inflight, [1, 1]);
1119
1120        book.admit(item(1, IO));
1121        assert_eq!(book.pending(), 1, "queued work counts against the cap immediately");
1122        book.admit(item(2, CPU));
1123        assert_eq!(book.pending(), 2);
1124        assert!(book.is_saturated());
1125        let dispatch = book.next(IO, Duration::ZERO).unwrap();
1126        assert_eq!(
1127            book.snapshot(),
1128            Snapshot {
1129                inflight: [1, 0],
1130                ready: [0, 1],
1131                pending: 2,
1132                resident: 2,
1133                evicting: 0,
1134                eviction_backlog: 0,
1135                queue_capacity: 2,
1136            }
1137        );
1138        book.complete(finish(dispatch), Duration::ZERO);
1139    }
1140
1141    #[test]
1142    fn three_classes_keep_separate_budgets() {
1143        let mut book: Scheduler<u64, &'static str, u64, 3> = Scheduler::new(Config {
1144            max_inflight: [1, 1, 1],
1145            max_pending: 8,
1146            max_resident: None,
1147            evict_after: Duration::from_secs(10),
1148            evict_iters: 8,
1149            queue_reserve: 8,
1150        });
1151        let now = Duration::ZERO;
1152        for class in 0..3u8 {
1153            book.admit(Admit { key: u64::from(class), class, expires_at: None, payload: "w" });
1154        }
1155        for class in 0..3u8 {
1156            let dispatch = book.next(class, now).expect("each class has its own budget");
1157            assert_eq!(dispatch.key, u64::from(class));
1158        }
1159        assert_eq!(book.snapshot().inflight, [1, 1, 1]);
1160        assert_eq!(book.check_invariants(), Ok(()));
1161    }
1162}