mongreldb_core/
node_governor.rs

1//! Node-level memory governor wiring (spec section 13.2, Stage 4B).
2//!
3//! Extends the Stage 1E [`MemoryGovernor`] across tablets and queries on one
4//! node: aggregates reservations, consumes OS-pressure inputs, and emits
5//! leader-move / admission actions when legal load still risks OOM.
6
7use std::collections::BTreeMap;
8
9use mongreldb_types::ids::TabletId;
10use serde::{Deserialize, Serialize};
11
12use crate::memory::{EscalationLevel, GovernorStats, MemoryGovernor};
13
14/// Inputs the node governor observes (spec §13.2).
15#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
16pub struct NodePressureInputs {
17    /// Physical memory bytes.
18    pub physical_memory_bytes: u64,
19    /// Configured maximum for the node.
20    pub configured_max_bytes: u64,
21    /// OS pressure signal (0.0..=1.0).
22    pub os_pressure: f64,
23    /// Aggregate cache hit rate (0.0..=1.0).
24    pub cache_hit_rate: f64,
25    /// Sum of query reservations.
26    pub query_reserved_bytes: u64,
27    /// Compaction backlog bytes.
28    pub compaction_backlog_bytes: u64,
29    /// Replication backlog bytes.
30    pub replication_backlog_bytes: u64,
31}
32
33impl Default for NodePressureInputs {
34    fn default() -> Self {
35        Self {
36            physical_memory_bytes: 16 * 1024 * 1024 * 1024,
37            configured_max_bytes: 12 * 1024 * 1024 * 1024,
38            os_pressure: 0.0,
39            cache_hit_rate: 1.0,
40            query_reserved_bytes: 0,
41            compaction_backlog_bytes: 0,
42            replication_backlog_bytes: 0,
43        }
44    }
45}
46
47/// Actions the governor may request (spec §13.2).
48#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
49pub enum GovernorAction {
50    /// Evict reclaimable caches.
51    EvictCaches,
52    /// Reduce query admission.
53    ReduceAdmission,
54    /// Spill analytics working sets.
55    SpillAnalytics,
56    /// Throttle compaction.
57    ThrottleCompaction,
58    /// Move tablet leaders off this node.
59    MoveTabletLeaders {
60        /// Suggested tablets to shed (may be empty = any).
61        tablets: Vec<TabletId>,
62    },
63    /// Reject oversized AI requests.
64    RejectOversizedAi,
65}
66
67/// Node-level governor: one [`MemoryGovernor`] plus tablet accounting.
68#[derive(Debug)]
69pub struct NodeMemoryGovernor {
70    /// Underlying governor.
71    pub governor: MemoryGovernor,
72    /// Per-tablet reserved bytes.
73    tablet_reserved: BTreeMap<TabletId, u64>,
74    /// Last computed actions.
75    last_actions: Vec<GovernorAction>,
76}
77
78impl NodeMemoryGovernor {
79    /// Wrap an existing governor.
80    pub fn new(governor: MemoryGovernor) -> Self {
81        Self {
82            governor,
83            tablet_reserved: BTreeMap::new(),
84            last_actions: Vec::new(),
85        }
86    }
87
88    /// Record a tablet reservation delta (positive = reserve, negative = release).
89    pub fn adjust_tablet(&mut self, tablet: TabletId, delta: i64) {
90        let entry = self.tablet_reserved.entry(tablet).or_insert(0);
91        if delta >= 0 {
92            *entry = entry.saturating_add(delta as u64);
93        } else {
94            *entry = entry.saturating_sub((-delta) as u64);
95        }
96        if *entry == 0 {
97            self.tablet_reserved.remove(&tablet);
98        }
99    }
100
101    /// Evaluate inputs and produce ordered actions (escalation ladder).
102    pub fn evaluate(&mut self, inputs: &NodePressureInputs) -> Vec<GovernorAction> {
103        let stats: GovernorStats = self.governor.stats();
104        let used = stats.total_used.saturating_add(inputs.query_reserved_bytes);
105        let max = inputs.configured_max_bytes.max(1);
106        let pressure = (used as f64 / max as f64).max(inputs.os_pressure);
107
108        let mut actions = Vec::new();
109        if pressure >= 0.70 || inputs.os_pressure >= 0.70 {
110            actions.push(GovernorAction::EvictCaches);
111        }
112        if pressure >= 0.80 {
113            actions.push(GovernorAction::ReduceAdmission);
114            actions.push(GovernorAction::SpillAnalytics);
115        }
116        if pressure >= 0.85 || inputs.compaction_backlog_bytes > max / 8 {
117            actions.push(GovernorAction::ThrottleCompaction);
118        }
119        if pressure >= 0.90 {
120            let tablets: Vec<TabletId> = self.tablet_reserved.keys().copied().collect();
121            actions.push(GovernorAction::MoveTabletLeaders { tablets });
122            actions.push(GovernorAction::RejectOversizedAi);
123        }
124
125        // Escalation level from the inner governor informs the same ladder.
126        match self.governor.escalation() {
127            EscalationLevel::None => {}
128            EscalationLevel::RejectLowPriority => {
129                if !actions.contains(&GovernorAction::ReduceAdmission) {
130                    actions.push(GovernorAction::ReduceAdmission);
131                }
132            }
133            EscalationLevel::EvictCaches | EscalationLevel::SpillOperators => {
134                if !actions.contains(&GovernorAction::EvictCaches) {
135                    actions.push(GovernorAction::EvictCaches);
136                }
137                if !actions.contains(&GovernorAction::SpillAnalytics) {
138                    actions.push(GovernorAction::SpillAnalytics);
139                }
140            }
141            EscalationLevel::ThrottleMaintenance => {
142                if !actions
143                    .iter()
144                    .any(|a| matches!(a, GovernorAction::MoveTabletLeaders { .. }))
145                {
146                    actions.push(GovernorAction::MoveTabletLeaders {
147                        tablets: self.tablet_reserved.keys().copied().collect(),
148                    });
149                }
150            }
151        }
152
153        self.last_actions = actions.clone();
154        actions
155    }
156
157    /// Last actions from [`evaluate`].
158    pub fn last_actions(&self) -> &[GovernorAction] {
159        &self.last_actions
160    }
161
162    /// Sum of per-tablet reservations.
163    pub fn tablet_reserved_bytes(&self) -> u64 {
164        self.tablet_reserved.values().sum()
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    use crate::memory::{GovernorConfig, MemoryGovernor};
172
173    fn gov(max: u64) -> NodeMemoryGovernor {
174        NodeMemoryGovernor::new(MemoryGovernor::new(GovernorConfig::new(max)).unwrap())
175    }
176
177    #[test]
178    fn high_pressure_emits_oom_prevention_actions() {
179        let mut g = gov(1_000_000);
180        let mut inputs = NodePressureInputs {
181            configured_max_bytes: 1_000_000,
182            query_reserved_bytes: 950_000,
183            os_pressure: 0.95,
184            ..NodePressureInputs::default()
185        };
186        let actions = g.evaluate(&inputs);
187        assert!(actions
188            .iter()
189            .any(|a| matches!(a, GovernorAction::EvictCaches)));
190        assert!(actions
191            .iter()
192            .any(|a| matches!(a, GovernorAction::RejectOversizedAi)));
193        assert!(actions
194            .iter()
195            .any(|a| matches!(a, GovernorAction::MoveTabletLeaders { .. })));
196        // Legal load still gets prevention actions — not silent OOM.
197        inputs.os_pressure = 0.5;
198        inputs.query_reserved_bytes = 1000;
199        let calm = g.evaluate(&inputs);
200        assert!(calm.len() < actions.len());
201    }
202
203    #[test]
204    fn tablet_accounting() {
205        let mut g = gov(10_000_000);
206        let t = TabletId::from_bytes([1; 16]);
207        g.adjust_tablet(t, 1000);
208        assert_eq!(g.tablet_reserved_bytes(), 1000);
209        g.adjust_tablet(t, -1000);
210        assert_eq!(g.tablet_reserved_bytes(), 0);
211    }
212}
mongreldb_core/node_governor.rs

mongreldb_core/
node_governor.rs
