Skip to main content

memory_budget/
config.rs

1//! Configurable knobs for [`MemoryBudget`](crate::MemoryBudget).
2
3use std::time::Duration;
4
5/// Default cadence at which the budget reapportions cache caps.
6pub(crate) const DEFAULT_TICK: Duration = Duration::from_secs(5);
7
8/// Upper hysteresis band — RSS above `target * SHRINK_RATIO` triggers
9/// a proportional shrink.
10pub(crate) const SHRINK_RATIO: f64 = 1.05;
11
12/// Lower hysteresis band — RSS below `target * GROW_RATIO` allows
13/// caches to grow back toward their preferred ceiling.
14pub(crate) const GROW_RATIO: f64 = 0.70;
15
16/// Ceiling on the auto-derived target when no `MEMORY_BUDGET_TARGET_MIB`
17/// is set: `min(0.8 * system_ram, 4096 MiB)`.
18const AUTO_TARGET_CEILING_MIB: u64 = 4096;
19
20/// Fraction of system RAM used when auto-deriving the target.
21const AUTO_TARGET_FRACTION: f64 = 0.8;
22
23/// Knobs that govern a [`MemoryBudget`](crate::MemoryBudget) instance.
24///
25/// Construct via [`BudgetConfig::from_env`] or build by hand for
26/// tests. Field values are absolute bytes (not MiB) so downstream
27/// code never has to second-guess the unit.
28#[derive(Clone, Debug)]
29pub struct BudgetConfig {
30    /// Soft RSS target in bytes.
31    target_bytes: u64,
32    /// Optional hard ceiling (process aborts on overrun).
33    hard_ceiling_bytes: Option<u64>,
34    /// Reapportionment cadence.
35    tick: Duration,
36}
37
38impl BudgetConfig {
39    /// Build a config explicitly. Useful for tests.
40    #[must_use]
41    pub fn new(target_bytes: u64) -> Self {
42        assert!(target_bytes > 0, "memory target must be greater than zero");
43        Self {
44            target_bytes,
45            hard_ceiling_bytes: None,
46            tick: DEFAULT_TICK,
47        }
48    }
49
50    /// Resolve a config from environment variables.
51    ///
52    /// Reads:
53    /// - `MEMORY_BUDGET_TARGET_MIB` (overrides the auto-derived target)
54    /// - `MEMORY_BUDGET_TICK_SECS` (default `5`)
55    /// - `MEMORY_BUDGET_HARD_CEILING_MIB` (optional)
56    ///
57    /// When `MEMORY_BUDGET_TARGET_MIB` is unset, the target is
58    /// `min(0.8 × system_ram, 4096 MiB)`. `system_ram_bytes` is the
59    /// installed RAM as seen by `sysinfo` and is required when the
60    /// target is derived automatically.
61    #[must_use]
62    pub fn from_env(system_ram_bytes: Option<u64>) -> Self {
63        let target_bytes = match read_mib_env("MEMORY_BUDGET_TARGET_MIB") {
64            Some(mib) => mib.saturating_mul(MIB),
65            None => auto_target_bytes(system_ram_bytes),
66        };
67        let hard_ceiling_bytes =
68            read_mib_env("MEMORY_BUDGET_HARD_CEILING_MIB").map(|mib| mib.saturating_mul(MIB));
69        let tick = read_secs_env("MEMORY_BUDGET_TICK_SECS")
70            .map(Duration::from_secs)
71            .unwrap_or(DEFAULT_TICK);
72        Self {
73            target_bytes,
74            hard_ceiling_bytes,
75            tick,
76        }
77    }
78
79    /// Soft RSS target in bytes.
80    #[must_use]
81    pub fn target_bytes(&self) -> u64 {
82        self.target_bytes
83    }
84
85    /// Optional hard abort ceiling.
86    #[must_use]
87    pub fn hard_ceiling_bytes(&self) -> Option<u64> {
88        self.hard_ceiling_bytes
89    }
90
91    /// Reapportionment cadence.
92    #[must_use]
93    pub fn tick(&self) -> Duration {
94        self.tick
95    }
96
97    /// Override the reapportionment cadence. Convenience for tests
98    /// that want to drive several ticks per second instead of one
99    /// every five.
100    #[must_use]
101    pub fn with_tick(mut self, tick: Duration) -> Self {
102        self.tick = tick;
103        self
104    }
105}
106
107const MIB: u64 = 1024 * 1024;
108
109fn auto_target_bytes(system_ram_bytes: Option<u64>) -> u64 {
110    let ram = system_ram_bytes
111        .expect("system RAM is required when MEMORY_BUDGET_TARGET_MIB is not configured");
112    assert!(ram > 0, "system RAM must be greater than zero");
113    let scaled = ((ram as f64) * AUTO_TARGET_FRACTION) as u64;
114    let target = scaled.min(AUTO_TARGET_CEILING_MIB.saturating_mul(MIB));
115    assert!(
116        target > 0,
117        "the derived memory target must be greater than zero"
118    );
119    target
120}
121
122fn read_mib_env(name: &str) -> Option<u64> {
123    let value = std::env::var(name).ok()?;
124    let mib = value
125        .trim()
126        .parse::<u64>()
127        .unwrap_or_else(|error| panic!("{name} must be an unsigned integer: {error}"));
128    assert!(mib > 0, "{name} must be greater than zero");
129    Some(mib)
130}
131
132fn read_secs_env(name: &str) -> Option<u64> {
133    let value = std::env::var(name).ok()?;
134    let seconds = value
135        .trim()
136        .parse::<u64>()
137        .unwrap_or_else(|error| panic!("{name} must be an unsigned integer: {error}"));
138    assert!(seconds > 0, "{name} must be greater than zero");
139    Some(seconds)
140}