molframe_core/execution/
memory.rs

1//! Shared accounting for bounded execution memory.
2
3use std::fmt;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::{Arc, Once, OnceLock};
6
7/// Smallest automatic execution budget, and the legacy fixed default: 100 MB.
8///
9/// The automatic default ([`MemoryBudget::auto`]) never falls below this. A
10/// caller who provisioned the machine is the only party who knows what the
11/// working set may be, so a budget larger or smaller than this is accepted
12/// without argument.
13pub const DEFAULT_MEMORY_BUDGET_BYTES: usize = 100_000_000;
14
15/// Environment variable that overrides the automatic budget (`4GiB`, `512MB`, bytes).
16pub const MEMORY_BUDGET_ENV: &str = "MOLFRAME_MEMORY_BUDGET";
17
18/// Largest automatic budget: 512 GiB. Explicit budgets are not capped.
19pub const MAX_AUTOMATIC_BUDGET_BYTES: usize = 512 << 30;
20
21/// Budget assumed when the platform cannot report its memory (1 GiB).
22const UNKNOWN_PLATFORM_BYTES: usize = 1 << 30;
23
24/// The text every memory refusal appends so the caller can act on it.
25pub const MEMORY_BUDGET_REMEDY: &str = "raise the memory budget: Python \
26ExecutionContext(memory_budget=...), Rust ExecutionContext::builder().memory_budget(MemoryBudget::new(..)?), \
27CLI --memory-budget, or environment variable MOLFRAME_MEMORY_BUDGET";
28
29/// A validated upper bound for memory retained by one execution context.
30#[derive(Clone, Copy, Debug, PartialEq, Eq)]
31pub struct MemoryBudget {
32    bytes: usize,
33}
34
35impl MemoryBudget {
36    /// Creates a non-zero budget.
37    ///
38    /// There is no compile-time maximum. A ceiling written into a library is a
39    /// ceiling nobody can raise, and a caller reading a hundred-gigabyte entry
40    /// needs a working set no constant here could have anticipated.
41    ///
42    /// # Errors
43    ///
44    /// Returns [`MemoryBudgetError::Zero`] when `bytes` is zero, because a
45    /// zero-byte budget cannot make forward progress.
46    pub const fn new(bytes: usize) -> Result<Self, MemoryBudgetError> {
47        if bytes == 0 {
48            return Err(MemoryBudgetError::Zero);
49        }
50        Ok(Self { bytes })
51    }
52
53    /// Returns the maximum number of retained bytes.
54    #[must_use]
55    pub const fn bytes(self) -> usize {
56        self.bytes
57    }
58
59    /// The budget this process uses when the caller sets none.
60    ///
61    /// `MOLFRAME_MEMORY_BUDGET` wins when it parses (see [`MemoryBudget::parse`]);
62    /// an unparsable value is reported once on stderr and ignored. Otherwise the
63    /// budget is [`MemoryBudget::for_physical_memory`] of the physical memory this
64    /// process may use: the container limit when one is set, available memory on
65    /// Linux, total memory on macOS. A platform that cannot report memory gets 1 GiB.
66    #[must_use]
67    pub fn auto() -> Self {
68        static DETECTED: OnceLock<MemoryBudget> = OnceLock::new();
69        if let Ok(text) = std::env::var(MEMORY_BUDGET_ENV) {
70            match Self::parse(&text) {
71                Ok(budget) => return budget,
72                Err(_) => warn_invalid_env(&text),
73            }
74        }
75        *DETECTED.get_or_init(|| match detect_physical_bytes() {
76            Some(bytes) => Self::for_physical_memory(bytes),
77            None => Self {
78                bytes: UNKNOWN_PLATFORM_BYTES,
79            },
80        })
81    }
82
83    /// Forty percent of `physical_bytes`, clamped to 100 MB through 512 GiB.
84    #[must_use]
85    pub fn for_physical_memory(physical_bytes: usize) -> Self {
86        let share = physical_bytes / 5 * 2;
87        Self {
88            bytes: share.clamp(DEFAULT_MEMORY_BUDGET_BYTES, MAX_AUTOMATIC_BUDGET_BYTES),
89        }
90    }
91
92    /// Parses a byte count such as `123456`, `512MB`, `4GiB` or `1.5G`.
93    ///
94    /// `KB`, `MB`, `GB`, `TB` are powers of 1000; `K`, `KiB` and so on, and a bare
95    /// `K`/`M`/`G`/`T`, are powers of 1024.
96    ///
97    /// # Errors
98    ///
99    /// Returns [`MemoryBudgetError::Zero`] for zero and
100    /// [`MemoryBudgetError::Unparsable`] for anything else that is not a size.
101    pub fn parse(text: &str) -> Result<Self, MemoryBudgetError> {
102        let trimmed = text.trim();
103        let split = match trimmed.find(|c: char| !(c.is_ascii_digit() || c == '.')) {
104            Some(split) => split,
105            None => trimmed.len(),
106        };
107        let (number, unit) = trimmed.split_at(split);
108        let multiplier: f64 = match unit.trim().to_ascii_lowercase().as_str() {
109            "" | "b" => 1.0,
110            "k" | "kib" => 1024.0,
111            "m" | "mib" => 1024.0 * 1024.0,
112            "g" | "gib" => 1024.0 * 1024.0 * 1024.0,
113            "t" | "tib" => 1024.0_f64.powi(4),
114            "kb" => 1e3,
115            "mb" => 1e6,
116            "gb" => 1e9,
117            "tb" => 1e12,
118            _ => return Err(MemoryBudgetError::Unparsable),
119        };
120        let value: f64 = number.parse().map_err(|_| MemoryBudgetError::Unparsable)?;
121        let bytes = value * multiplier;
122        let ceiling = crate::numeric::usize_to_f64(usize::MAX);
123        if !bytes.is_finite() || bytes >= ceiling {
124            return Err(MemoryBudgetError::Unparsable);
125        }
126        // Truncation is the intent: a fractional byte is not a byte.
127        Self::new(crate::numeric::f64_to_usize(bytes))
128    }
129}
130
131impl Default for MemoryBudget {
132    /// The machine-scaled budget of [`MemoryBudget::auto`].
133    fn default() -> Self {
134        Self::auto()
135    }
136}
137
138fn warn_invalid_env(text: &str) {
139    static WARNED: Once = Once::new();
140    WARNED.call_once(|| {
141        eprintln!(
142            "molframe: ignoring {MEMORY_BUDGET_ENV}={text:?}: expected a positive size such as 4GiB or 512MB"
143        );
144    });
145}
146
147/// Physical bytes this process may use, when the platform says so cheaply.
148fn detect_physical_bytes() -> Option<usize> {
149    let limit = cgroup_limit();
150    let physical = if cfg!(target_os = "linux") {
151        std::fs::read_to_string("/proc/meminfo")
152            .ok()
153            .and_then(|text| {
154                meminfo_kib(&text, "MemAvailable:").or_else(|| meminfo_kib(&text, "MemTotal:"))
155            })
156            .map(|kib| kib.saturating_mul(1024))
157    } else if cfg!(target_os = "macos") {
158        std::process::Command::new("sysctl")
159            .args(["-n", "hw.memsize"])
160            .output()
161            .ok()
162            .and_then(|out| String::from_utf8(out.stdout).ok())
163            .and_then(|text| text.trim().parse::<usize>().ok())
164    } else {
165        None
166    };
167    match (physical, limit) {
168        (Some(a), Some(b)) => Some(a.min(b)),
169        (a, b) => a.or(b),
170    }
171}
172
173fn meminfo_kib(text: &str, key: &str) -> Option<usize> {
174    text.lines()
175        .find_map(|line| line.strip_prefix(key))
176        .and_then(|rest| rest.split_whitespace().next())
177        .and_then(|value| value.parse().ok())
178}
179
180/// The container memory limit (cgroup v2, then v1), when one is set.
181fn cgroup_limit() -> Option<usize> {
182    [
183        "/sys/fs/cgroup/memory.max",
184        "/sys/fs/cgroup/memory/memory.limit_in_bytes",
185    ]
186    .iter()
187    .find_map(|path| std::fs::read_to_string(path).ok())
188    .and_then(|text| text.trim().parse::<u64>().ok())
189    // cgroup v1 reports "no limit" as a number near 2^63. The comparison is
190    // done in u64 so a 32-bit target does not overflow the shift.
191    .filter(|limit| *limit < (1_u64 << 60))
192    .and_then(|limit| usize::try_from(limit).ok())
193}
194
195/// Failure to construct or reserve an execution-memory budget.
196#[derive(Clone, Copy, Debug, PartialEq, Eq)]
197pub enum MemoryBudgetError {
198    /// A zero-byte budget cannot make forward progress.
199    Zero,
200    /// The text is not a positive size such as `4GiB`.
201    Unparsable,
202    /// The context has insufficient unreserved capacity.
203    Exhausted {
204        /// Bytes requested by the operation.
205        requested: usize,
206        /// Bytes still available at the time of the attempt.
207        available: usize,
208    },
209}
210
211impl fmt::Display for MemoryBudgetError {
212    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
213        match self {
214            Self::Zero => formatter.write_str("execution memory budget must be non-zero"),
215            Self::Unparsable => {
216                formatter.write_str("memory budget must be a positive size such as 4GiB or 512MB")
217            }
218            Self::Exhausted {
219                requested,
220                available,
221            } => write!(
222                formatter,
223                "execution requested {requested} bytes with {available} bytes available; \
224                 {MEMORY_BUDGET_REMEDY}"
225            ),
226        }
227    }
228}
229
230impl std::error::Error for MemoryBudgetError {}
231
232#[derive(Debug)]
233pub(crate) struct MemoryAccount {
234    budget: MemoryBudget,
235    reserved: AtomicUsize,
236    peak_reserved: AtomicUsize,
237    live_batches: AtomicUsize,
238    peak_live_batches: AtomicUsize,
239}
240
241impl MemoryAccount {
242    pub(crate) fn new(budget: MemoryBudget) -> Self {
243        Self {
244            budget,
245            reserved: AtomicUsize::new(0),
246            peak_reserved: AtomicUsize::new(0),
247            live_batches: AtomicUsize::new(0),
248            peak_live_batches: AtomicUsize::new(0),
249        }
250    }
251
252    pub(crate) const fn budget(&self) -> MemoryBudget {
253        self.budget
254    }
255
256    pub(crate) fn reserve(
257        account: &Arc<Self>,
258        bytes: usize,
259    ) -> Result<MemoryReservation, MemoryBudgetError> {
260        let mut current = account.reserved.load(Ordering::Acquire);
261        loop {
262            let Some(next) = current.checked_add(bytes) else {
263                return Err(MemoryBudgetError::Exhausted {
264                    requested: bytes,
265                    available: account.budget.bytes().saturating_sub(current),
266                });
267            };
268            if next > account.budget.bytes() {
269                return Err(MemoryBudgetError::Exhausted {
270                    requested: bytes,
271                    available: account.budget.bytes().saturating_sub(current),
272                });
273            }
274            match account.reserved.compare_exchange_weak(
275                current,
276                next,
277                Ordering::AcqRel,
278                Ordering::Acquire,
279            ) {
280                Ok(_) => {
281                    account.peak_reserved.fetch_max(next, Ordering::AcqRel);
282                    return Ok(MemoryReservation {
283                        account: Arc::clone(account),
284                        bytes,
285                    });
286                }
287                Err(observed) => current = observed,
288            }
289        }
290    }
291
292    pub(crate) fn reserved_bytes(&self) -> usize {
293        self.reserved.load(Ordering::Acquire)
294    }
295
296    pub(crate) fn peak_reserved_bytes(&self) -> usize {
297        self.peak_reserved.load(Ordering::Acquire)
298    }
299
300    pub(crate) fn live_batches(&self) -> usize {
301        self.live_batches.load(Ordering::Acquire)
302    }
303
304    pub(crate) fn peak_live_batches(&self) -> usize {
305        self.peak_live_batches.load(Ordering::Acquire)
306    }
307
308    fn begin_batch(account: &Arc<Self>) -> BatchGuard {
309        let live = account.live_batches.fetch_add(1, Ordering::AcqRel) + 1;
310        account.peak_live_batches.fetch_max(live, Ordering::AcqRel);
311        BatchGuard {
312            account: Arc::clone(account),
313        }
314    }
315}
316
317/// Exclusive accounting for bytes retained by an operation.
318///
319/// Dropping the value returns its capacity to the context. Reservations cannot
320/// be cloned, preventing one charge from representing multiple owners.
321#[derive(Debug)]
322pub struct MemoryReservation {
323    account: Arc<MemoryAccount>,
324    bytes: usize,
325}
326
327impl MemoryReservation {
328    /// Reports whether this reservation belongs to the supplied execution account.
329    #[must_use]
330    pub fn belongs_to(&self, context: &super::ExecutionContext) -> bool {
331        Arc::ptr_eq(&self.account, &context.memory)
332    }
333
334    /// Returns the charged byte count.
335    #[must_use]
336    pub const fn bytes(&self) -> usize {
337        self.bytes
338    }
339
340    /// Reduces this reservation to the exact retained byte count.
341    pub fn shrink_to(&mut self, bytes: usize) {
342        if bytes >= self.bytes {
343            return;
344        }
345        let released = self.bytes - bytes;
346        self.bytes = bytes;
347        self.account.reserved.fetch_sub(released, Ordering::AcqRel);
348    }
349
350    /// Extends this reservation before its owner grows an allocation.
351    ///
352    /// # Errors
353    ///
354    /// Returns [`MemoryBudgetError::Exhausted`] without changing this
355    /// reservation when the additional bytes do not fit.
356    pub fn try_grow(&mut self, additional: usize) -> Result<(), MemoryBudgetError> {
357        if additional == 0 {
358            return Ok(());
359        }
360        let mut extra = MemoryAccount::reserve(&self.account, additional)?;
361        let Some(bytes) = self.bytes.checked_add(additional) else {
362            return Err(MemoryBudgetError::Exhausted {
363                requested: additional,
364                available: 0,
365            });
366        };
367        self.bytes = bytes;
368        extra.bytes = 0;
369        Ok(())
370    }
371
372    pub(crate) fn begin_batch(&self) -> BatchGuard {
373        MemoryAccount::begin_batch(&self.account)
374    }
375}
376
377impl Drop for MemoryReservation {
378    fn drop(&mut self) {
379        self.account
380            .reserved
381            .fetch_sub(self.bytes, Ordering::AcqRel);
382    }
383}
384
385/// Exactly-once accounting for one live batch lease.
386#[derive(Debug)]
387pub(crate) struct BatchGuard {
388    account: Arc<MemoryAccount>,
389}
390
391impl Drop for BatchGuard {
392    fn drop(&mut self) {
393        self.account.live_batches.fetch_sub(1, Ordering::AcqRel);
394    }
395}
396
397#[cfg(test)]
398#[path = "memory_tests.rs"]
399mod tests;
400
401crate::diagnostic_from!(
402    MemoryBudgetError,
403    |_error| crate::Code::E7001,
404    |diagnostic, error| match error {
405        MemoryBudgetError::Zero => diagnostic.with_context("budget", "zero"),
406        MemoryBudgetError::Unparsable => diagnostic.with_context("budget", "unparsable"),
407        MemoryBudgetError::Exhausted {
408            requested,
409            available,
410        } => diagnostic
411            .with_context("requested_bytes", requested.to_string())
412            .with_context("available_bytes", available.to_string()),
413    }
414);
molframe_core/execution/memory.rs

molframe_core/execution/
memory.rs
