memory_budget/budget/memory/
accessors.rs

1//! Read-only accessors and the lease API on [`MemoryBudget`].
2
3use std::sync::Arc;
4use std::sync::atomic::Ordering;
5use std::time::Duration;
6
7use super::handle::MemoryBudget;
8use crate::ReservationError;
9use crate::lease::LeaseGuard;
10
11const EXECUTOR_SPILL_TARGET_DIVISOR: u64 = 16;
12const MIN_EXECUTOR_SPILL_RUN_BUFFER_BYTES: u64 = 1024 * 1024;
13const MAX_EXECUTOR_SPILL_RUN_BUFFER_BYTES: u64 = 32 * 1024 * 1024;
14
15const RESULT_CHANNEL_TARGET_DIVISOR: u64 = 32;
16const MIN_RESULT_CHANNEL_BYTES: u64 = 256 * 1024;
17const MAX_RESULT_CHANNEL_BYTES: u64 = 8 * 1024 * 1024;
18
19impl MemoryBudget {
20    /// Observe current process/allocator usage without running a cache tick.
21    ///
22    /// This is a measurement, not an allocation reservation. The maximum of
23    /// process usage, allocator residency and live heap bytes protects against
24    /// both compressed pages and allocations outside the configured allocator.
25    #[must_use]
26    pub fn observed_usage_bytes(&self) -> u64 {
27        self.state
28            .jemalloc
29            .read()
30            .observed_usage_bytes(self.state.rss.read_rss_bytes())
31    }
32
33    /// Reserve `bytes`, panicking if outstanding leases would exceed the target.
34    ///
35    /// Use [`Self::try_lease`] when exhaustion is an expected refusal. The
36    /// reservation must precede allocation and follow its ownership to release.
37    /// This bounds outstanding reservations, not unregistered allocations.
38    ///
39    /// # Panics
40    ///
41    /// Panics when the aggregate reservation would exceed the target.
42    #[must_use = "dropping the guard immediately releases the lease"]
43    pub fn lease(&self, bytes: u64) -> LeaseGuard {
44        self.try_lease(bytes)
45            .unwrap_or_else(|error| panic!("{error}"))
46    }
47
48    /// Reserve bytes atomically without exceeding the aggregate lease target.
49    ///
50    /// Registered cache ceilings are reduced by active leases on the next
51    /// coordination tick. Existing allocations are not synchronously evicted;
52    /// callers must account for cache/baseline ownership separately before
53    /// claiming an enforced whole-process allocation bound.
54    ///
55    /// # Errors
56    ///
57    /// Returns [`ReservationError`] without changing accounting when the
58    /// request cannot fit. No payload is allocated and no waiter is queued.
59    pub fn try_lease(&self, bytes: u64) -> Result<LeaseGuard, ReservationError> {
60        let capacity = self.target_bytes();
61        let mut reserved = self.state.active_lease_bytes.load(Ordering::Relaxed);
62        loop {
63            let Some(next) = reserved.checked_add(bytes).filter(|next| *next <= capacity) else {
64                return Err(ReservationError::new(bytes, reserved, capacity));
65            };
66            match self.state.active_lease_bytes.compare_exchange_weak(
67                reserved,
68                next,
69                Ordering::Relaxed,
70                Ordering::Relaxed,
71            ) {
72                Ok(_) => return Ok(LeaseGuard::new(bytes, Arc::clone(&self.state))),
73                Err(observed) => reserved = observed,
74            }
75        }
76    }
77
78    /// Soft RSS target in bytes, before any lease deductions.
79    #[must_use]
80    pub fn target_bytes(&self) -> u64 {
81        self.state.config.target_bytes()
82    }
83
84    /// Effective target after subtracting admitted active leases.
85    #[must_use]
86    pub fn effective_target_bytes(&self) -> u64 {
87        self.state
88            .config
89            .target_bytes()
90            .saturating_sub(self.state.active_lease_bytes.load(Ordering::Relaxed))
91    }
92
93    /// Per-operator external-sort run buffer for executor spill paths.
94    ///
95    /// Derived from the effective process target, so active leases shrink
96    /// executor spill buffers before an operator starts rather than letting
97    /// every operator size itself against the whole process budget.
98    ///
99    /// # Panics
100    ///
101    /// Panics if the bounded spill-buffer size cannot fit in `usize`.
102    #[must_use]
103    pub fn executor_spill_run_buffer_bytes(&self) -> usize {
104        let bytes = (self.effective_target_bytes() / EXECUTOR_SPILL_TARGET_DIVISOR).clamp(
105            MIN_EXECUTOR_SPILL_RUN_BUFFER_BYTES,
106            MAX_EXECUTOR_SPILL_RUN_BUFFER_BYTES,
107        );
108        usize::try_from(bytes).expect("executor spill run buffer fits usize")
109    }
110
111    /// Byte budget for one result stream's producer→consumer channel
112    /// (BME-22): the ceiling on encoded result bytes in flight between
113    /// the encoder and its consumer/spool tee. A **byte** bound, not a
114    /// frame count — frames vary in size by orders of magnitude, so
115    /// counting them would bound the wrong quantity (BME-12's lesson).
116    ///
117    /// Derived from the effective process target so active leases shrink
118    /// in-flight result buffering the same way they shrink spill
119    /// buffers.
120    ///
121    /// # Panics
122    ///
123    /// Panics if the bounded channel size cannot fit in `usize`.
124    #[must_use]
125    pub fn result_channel_bytes(&self) -> usize {
126        let bytes = (self.effective_target_bytes() / RESULT_CHANNEL_TARGET_DIVISOR)
127            .clamp(MIN_RESULT_CHANNEL_BYTES, MAX_RESULT_CHANNEL_BYTES);
128        usize::try_from(bytes).expect("result channel size fits usize")
129    }
130
131    /// Tick cadence used by the spawned background loop.
132    #[must_use]
133    pub fn tick_interval(&self) -> Duration {
134        self.state.config.tick()
135    }
136
137    /// Optional hard ceiling. When set, the budget aborts the
138    /// process if RSS exceeds this value — a backstop for runaway
139    /// growth in production.
140    #[must_use]
141    pub fn hard_ceiling_bytes(&self) -> Option<u64> {
142        self.state.config.hard_ceiling_bytes()
143    }
144}
memory_budget/budget/memory/accessors.rs

memory_budget/budget/memory/
accessors.rs
