Skip to main content

nodedb_mem/
engine_limits.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Total per-engine byte limit map.
4//!
5//! `EngineLimits` holds exactly one byte limit for every `EngineId` —
6//! absence is unrepresentable. A limit of 0 is a valid config: the engine
7//! rejects every non-zero reservation, it is never a missing-entry error.
8
9use crate::engine::EngineId;
10
11/// A byte limit for every `EngineId`, indexed by `EngineId::index()`.
12#[derive(Debug, Clone)]
13pub struct EngineLimits([usize; EngineId::COUNT]);
14
15impl EngineLimits {
16    /// Every engine limited to zero bytes.
17    pub const fn zeroed() -> Self {
18        Self([0; EngineId::COUNT])
19    }
20
21    /// Every engine limited to the same number of bytes.
22    pub const fn uniform(bytes: usize) -> Self {
23        Self([bytes; EngineId::COUNT])
24    }
25
26    /// Consuming builder: set `engine`'s limit to `bytes`.
27    pub fn with(mut self, engine: EngineId, bytes: usize) -> Self {
28        self.set(engine, bytes);
29        self
30    }
31
32    /// Set `engine`'s limit to `bytes`.
33    pub fn set(&mut self, engine: EngineId, bytes: usize) {
34        self.0[engine.index()] = bytes;
35    }
36
37    /// The byte limit configured for `engine`.
38    pub fn get(&self, engine: EngineId) -> usize {
39        self.0[engine.index()]
40    }
41
42    /// Sum of every engine's limit.
43    pub fn total(&self) -> usize {
44        self.0.iter().sum()
45    }
46
47    /// Iterate `(engine, limit)` pairs in `EngineId::ALL` order.
48    pub fn iter(&self) -> impl Iterator<Item = (EngineId, usize)> + '_ {
49        EngineId::ALL
50            .iter()
51            .map(|&engine| (engine, self.get(engine)))
52    }
53
54    /// Borrow the raw per-engine array, indexed by `EngineId::index()`.
55    pub(crate) fn as_array(&self) -> &[usize; EngineId::COUNT] {
56        &self.0
57    }
58}
59
60impl FromIterator<(EngineId, usize)> for EngineLimits {
61    /// Builds a total map from a partial iterator. An engine not yielded
62    /// keeps its zero default rather than being absent.
63    fn from_iter<T: IntoIterator<Item = (EngineId, usize)>>(iter: T) -> Self {
64        let mut limits = Self::zeroed();
65        for (engine, bytes) in iter {
66            limits.set(engine, bytes);
67        }
68        limits
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn zeroed_has_no_limit() {
78        let limits = EngineLimits::zeroed();
79        for &engine in EngineId::ALL {
80            assert_eq!(limits.get(engine), 0);
81        }
82        assert_eq!(limits.total(), 0);
83    }
84
85    #[test]
86    fn uniform_applies_to_every_engine() {
87        let limits = EngineLimits::uniform(1024);
88        for &engine in EngineId::ALL {
89            assert_eq!(limits.get(engine), 1024);
90        }
91        assert_eq!(limits.total(), 1024 * EngineId::COUNT);
92    }
93
94    #[test]
95    fn with_sets_a_single_engine() {
96        let limits = EngineLimits::zeroed()
97            .with(EngineId::Vector, 100)
98            .with(EngineId::Query, 200);
99        assert_eq!(limits.get(EngineId::Vector), 100);
100        assert_eq!(limits.get(EngineId::Query), 200);
101        assert_eq!(limits.get(EngineId::Kv), 0);
102        assert_eq!(limits.total(), 300);
103    }
104
105    #[test]
106    fn set_mutates_in_place() {
107        let mut limits = EngineLimits::zeroed();
108        limits.set(EngineId::Crdt, 42);
109        assert_eq!(limits.get(EngineId::Crdt), 42);
110    }
111
112    #[test]
113    fn iter_covers_every_engine_in_all_order() {
114        let limits = EngineLimits::zeroed().with(EngineId::Fts, 7);
115        let collected: Vec<_> = limits.iter().collect();
116        assert_eq!(collected.len(), EngineId::COUNT);
117        assert_eq!(
118            collected.iter().map(|(e, _)| *e).collect::<Vec<_>>(),
119            EngineId::ALL.to_vec()
120        );
121        assert_eq!(
122            collected
123                .iter()
124                .find(|(e, _)| *e == EngineId::Fts)
125                .unwrap()
126                .1,
127            7
128        );
129    }
130
131    #[test]
132    fn from_iter_leaves_unlisted_engines_zero() {
133        let limits: EngineLimits = [(EngineId::Vector, 10), (EngineId::Kv, 20)]
134            .into_iter()
135            .collect();
136        assert_eq!(limits.get(EngineId::Vector), 10);
137        assert_eq!(limits.get(EngineId::Kv), 20);
138        assert_eq!(limits.get(EngineId::Graph), 0);
139    }
140}