wowlab_sentinel/scheduler/
backlog.rs

1//! Canonical in-memory scheduler backlog accounting.
2
3use wowlab_common::NodePublicKey;
4use wowlab_types::sim::FastMap;
5
6use super::runtime::JobRuntimeStore;
7
8#[derive(Debug, Default)]
9pub(super) struct NodeBacklogs(FastMap<NodePublicKey, usize>);
10
11impl NodeBacklogs {
12    pub(super) fn from_runtimes(runtimes: &JobRuntimeStore) -> Self {
13        let mut backlogs = FastMap::default();
14
15        for runtime in runtimes {
16            for claim in runtime.in_flight_claims() {
17                if let Ok(public_key) = claim.node_public_key.parse::<NodePublicKey>() {
18                    *backlogs.entry(public_key).or_insert(0) += 1;
19                }
20            }
21        }
22
23        Self(backlogs)
24    }
25
26    pub(super) fn available(&self, public_key: &NodePublicKey, capacity: usize) -> Option<usize> {
27        let backlog = self.0.get(public_key).copied().unwrap_or(0);
28
29        (backlog < capacity).then_some(capacity - backlog)
30    }
31
32    pub(super) fn increment(&mut self, public_key: &NodePublicKey) {
33        *self.0.entry(public_key.clone()).or_insert(0) += 1;
34    }
35}
36
37pub(crate) fn pending_claim_count(runtimes: &JobRuntimeStore) -> i64 {
38    runtimes
39        .iter()
40        .map(|job| i64::try_from(job.in_flight_len()).unwrap_or(i64::MAX))
41        .sum()
42}
43
44#[cfg(test)]
45mod tests {
46    use googletest::prelude::*;
47
48    use super::*;
49
50    #[gtest]
51    fn empty_backlog_exposes_full_capacity() -> Result<()> {
52        let backlogs = NodeBacklogs::default();
53
54        verify_eq!(
55            backlogs.available(&NodePublicKey::from_bytes([1; 32]), 4),
56            Some(4)
57        )?;
58
59        Ok(())
60    }
61
62    #[gtest]
63    fn increment_enforces_backpressure_at_capacity() -> Result<()> {
64        let key = NodePublicKey::from_bytes([2; 32]);
65        let mut backlogs = NodeBacklogs::default();
66
67        backlogs.increment(&key);
68        verify_eq!(backlogs.available(&key, 2), Some(1))?;
69        backlogs.increment(&key);
70        verify_eq!(backlogs.available(&key, 2), None)?;
71
72        Ok(())
73    }
74}
wowlab_sentinel/scheduler/backlog.rs

wowlab_sentinel/scheduler/
backlog.rs
