nodedb_mem/governor/
core.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! [`MemoryGovernor`] struct definition, construction, and basic accessors.
4
5use std::collections::HashMap;
6use std::sync::{Arc, RwLock};
7
8use nodedb_types::{DatabaseId, TenantId};
9
10use super::config::GovernorConfig;
11use super::global_counter::GlobalCounter;
12use crate::budget::Budget;
13use crate::engine::EngineId;
14use crate::error::Result;
15use crate::pressure::PressureThresholds;
16use crate::scoped_budget::ScopedBudget;
17
18/// The central memory governor.
19///
20/// Thread-safe: global, database, and tenant counters use atomics.
21/// The budget map itself is behind an `RwLock`; reads (common) take a shared
22/// lock, writes (rare — only when quotas change) take an exclusive lock.
23#[derive(Debug)]
24pub struct MemoryGovernor {
25    /// Per-engine budgets, one for every `EngineId`, indexed by `EngineId::index()`.
26    pub(super) budgets: [Budget; EngineId::COUNT],
27
28    /// Shared global counter. Held by both the governor and every live token.
29    pub(crate) global_counter: Arc<GlobalCounter>,
30
31    /// Global ceiling in bytes.
32    pub(super) global_ceiling: usize,
33
34    /// Pressure thresholds for graduated backpressure.
35    pub(super) thresholds: PressureThresholds,
36
37    /// Per-database budget map. Keyed by `DatabaseId`. Populated lazily via
38    /// `set_database_budget`; databases without an entry are uncapped.
39    pub(super) database_budgets: RwLock<HashMap<DatabaseId, ScopedBudget>>,
40
41    /// Per-tenant budget map. Keyed by `(DatabaseId, TenantId)`. Populated
42    /// lazily via `set_tenant_budget`.
43    pub(super) tenant_budgets: RwLock<HashMap<(DatabaseId, TenantId), ScopedBudget>>,
44}
45
46impl MemoryGovernor {
47    /// Create a new governor with the given configuration.
48    pub fn new(config: GovernorConfig) -> Result<Self> {
49        config.validate()?;
50
51        let limits = config.engine_limits.as_array();
52        let budgets: [Budget; EngineId::COUNT] = std::array::from_fn(|i| Budget::new(limits[i]));
53
54        let global_counter = Arc::new(GlobalCounter::new(config.global_ceiling));
55
56        Ok(Self {
57            budgets,
58            global_counter,
59            global_ceiling: config.global_ceiling,
60            thresholds: PressureThresholds::default(),
61            database_budgets: RwLock::new(HashMap::new()),
62            tenant_budgets: RwLock::new(HashMap::new()),
63        })
64    }
65
66    /// Get the budget for a specific engine.
67    pub fn budget(&self, engine: EngineId) -> &Budget {
68        &self.budgets[engine.index()]
69    }
70
71    /// Get the global ceiling.
72    pub fn global_ceiling(&self) -> usize {
73        self.global_ceiling
74    }
75}
nodedb_mem/governor/core.rs

nodedb_mem/governor/
core.rs
