mz_compute/
memory_limiter.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Utilities to limit memory usage.
11//!
12//! In the context of this module, "memory" refers to the sum of physical memory and swap space .
13//! Other parts of the code usually don't include swap space when talking about "memory".
14
15use std::sync::Mutex;
16use std::time::{Duration, Instant};
17
18use anyhow::Context;
19use mz_compute_types::dyncfgs::{
20    MEMORY_LIMITER_BURST_FACTOR, MEMORY_LIMITER_INTERVAL, MEMORY_LIMITER_USAGE_BIAS,
21};
22use mz_dyncfg::ConfigSet;
23use mz_ore::cast::{CastFrom, CastLossy};
24use mz_ore::metric;
25use mz_ore::metrics::{MetricsRegistry, UIntGauge};
26use prometheus::Histogram;
27use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
28use tracing::{debug, error, info, warn};
29
30/// A handle to the active memory limiter.
31///
32/// The limiter is initialized by a call to [`start_limiter`]. It spawns a process-global task that
33/// runs for the lifetime of the process.
34static LIMITER: Mutex<Option<Limiter>> = Mutex::new(None);
35
36/// Start the process-global memory limiter.
37///
38/// # Panics
39///
40/// Panics if the limiter was already started previously.
41pub fn start_limiter(memory_limit: usize, metrics_registry: &MetricsRegistry) {
42    let mut limiter = LIMITER.lock().expect("poisoned");
43
44    if limiter.is_some() {
45        panic!("memory limiter is already running");
46    }
47
48    let metrics = LimiterMetrics::new(metrics_registry);
49    let (config_tx, config_rx) = mpsc::unbounded_channel();
50
51    mz_ore::task::spawn(|| "memory-limiter", LimiterTask::run(config_rx, metrics));
52
53    *limiter = Some(Limiter {
54        base_memory_limit: memory_limit,
55        effective_memory_limit: memory_limit,
56        config_tx,
57    });
58}
59
60/// Apply the given configuration to the active limiter.
61pub fn apply_limiter_config(config: &ConfigSet) {
62    if let Some(limiter) = LIMITER.lock().expect("poisoned").as_mut() {
63        limiter.apply_config(config);
64    }
65}
66
67/// Get the current effective memory limit.
68pub fn get_memory_limit() -> Option<usize> {
69    let limiter = LIMITER.lock().expect("poisoned");
70    limiter.as_ref().map(|l| l.effective_memory_limit)
71}
72
73/// A handle to a running memory limiter task.
74struct Limiter {
75    /// The base process memory limit.
76    base_memory_limit: usize,
77    /// The effective memory limit, obtained by applying dyncfgs to the base limit.
78    effective_memory_limit: usize,
79    /// A sender for limiter configuration updates.
80    config_tx: UnboundedSender<LimiterConfig>,
81}
82
83impl Limiter {
84    /// Apply the given configuration to the limiter.
85    fn apply_config(&mut self, config: &ConfigSet) {
86        let mut interval = MEMORY_LIMITER_INTERVAL.get(config);
87        // A zero duration means the limiter is disabled. Translate that into an ~infinite duration
88        // so the limiter doesn't have to worry about the special case.
89        if interval.is_zero() {
90            interval = Duration::MAX;
91        }
92
93        let memory_limit =
94            f64::cast_lossy(self.base_memory_limit) * MEMORY_LIMITER_USAGE_BIAS.get(config);
95        let memory_limit = usize::cast_lossy(memory_limit);
96
97        let burst_budget = f64::cast_lossy(memory_limit) * MEMORY_LIMITER_BURST_FACTOR.get(config);
98        let burst_budget = usize::cast_lossy(burst_budget);
99
100        self.effective_memory_limit = memory_limit;
101
102        self.config_tx
103            .send(LimiterConfig {
104                interval,
105                memory_limit,
106                burst_budget,
107            })
108            .expect("limiter task never shuts down");
109    }
110}
111
112/// Configuration for an memory limiter task.
113#[derive(Clone, Copy, Debug, PartialEq, Eq)]
114struct LimiterConfig {
115    /// The interval at which memory usage is checked against the memory limit.
116    interval: Duration,
117    /// The memory limit.
118    memory_limit: usize,
119    /// Budget to allow memory usage above the memory limit, in byte-seconds.
120    burst_budget: usize,
121}
122
123impl LimiterConfig {
124    /// Return a config that disables the memory limiter.
125    fn disabled() -> Self {
126        Self {
127            interval: Duration::MAX,
128            memory_limit: 0,
129            burst_budget: 0,
130        }
131    }
132}
133
134/// A task that enforces configured memory limits.
135///
136/// The task operates by performing limit checks at a configured interval. For each check it
137/// obtains the current memory utilization from proc stats. It then compares the utilization against
138/// the configured memory limit, and if it exceeds the limit, reduces the burst budget by the amount
139/// of memory utilization that exceeds the limit. If the burst budget is exhausted, the limiter
140/// terminates the process.
141struct LimiterTask {
142    /// The current limiter configuration.
143    config: LimiterConfig,
144    /// The amount of burst budget remaining.
145    burst_budget_remaining: usize,
146    /// The time of the last check.
147    last_check: Instant,
148    /// Metrics tracked by the limiter task.
149    metrics: LimiterMetrics,
150}
151
152impl LimiterTask {
153    async fn run(mut config_rx: UnboundedReceiver<LimiterConfig>, metrics: LimiterMetrics) {
154        info!("running memory limiter task");
155
156        let mut task = Self {
157            config: LimiterConfig::disabled(),
158            burst_budget_remaining: 0,
159            last_check: Instant::now(),
160            metrics,
161        };
162
163        loop {
164            tokio::select! {
165                _ = task.tick() => {
166                    let start = Instant::now();
167
168                    if let Err(err) = task.check() {
169                        error!("memory limit check failed: {err}");
170                    }
171
172                    let elapsed = start.elapsed();
173                    task.metrics.duration.observe(elapsed.as_secs_f64());
174                }
175                Some(config) = config_rx.recv() => task.apply_config(config),
176            }
177        }
178    }
179
180    /// Wait until the next check time.
181    async fn tick(&self) {
182        let elapsed = self.last_check.elapsed();
183        let duration = self.config.interval.saturating_sub(elapsed);
184        tokio::time::sleep(duration).await
185    }
186
187    fn current_utilization() -> anyhow::Result<ProcStatus> {
188        match ProcStatus::from_proc() {
189            Ok(status) => Ok(status),
190            #[cfg(target_os = "linux")]
191            Err(err) => {
192                error!("failed to read /proc/self/status: {err}");
193                Err(err)
194            }
195            #[cfg(not(target_os = "linux"))]
196            Err(_err) => Ok(ProcStatus {
197                vm_rss: 0,
198                vm_swap: 0,
199            }),
200        }
201    }
202
203    /// Perform a memory usage check, terminating the process if the configured limits are exceeded.
204    fn check(&mut self) -> Result<(), anyhow::Error> {
205        debug!("checking memory limits");
206
207        let ProcStatus { vm_rss, vm_swap } = Self::current_utilization()?;
208
209        let memory_limit = self.config.memory_limit;
210        let burst_budget_remaining = self.burst_budget_remaining;
211
212        let memory_usage = vm_rss + vm_swap;
213
214        debug!(
215            memory_usage,
216            memory_limit, burst_budget_remaining, vm_rss, vm_swap, "memory utilization check",
217        );
218
219        self.metrics.vm_rss.set(u64::cast_from(vm_rss));
220        self.metrics.vm_swap.set(u64::cast_from(vm_swap));
221        self.metrics.memory_usage.set(u64::cast_from(memory_usage));
222        self.metrics
223            .burst_budget
224            .set(u64::cast_from(burst_budget_remaining));
225
226        if memory_usage > memory_limit {
227            // Calculate excess usage in byte-seconds.
228            let elapsed = self.last_check.elapsed().as_secs_f64();
229            let excess = memory_usage - memory_limit;
230            let excess_bs = usize::cast_lossy(f64::cast_lossy(excess) * elapsed);
231
232            if burst_budget_remaining >= excess_bs {
233                self.burst_budget_remaining -= excess_bs;
234            } else {
235                // Burst budget exhausted, terminate the process.
236                warn!(
237                    memory_usage,
238                    memory_limit, "memory utilization exceeded configured limits",
239                );
240                // We terminate with a recognizable exit code so the orchestrator knows the
241                // termination was caused by exceeding memory limits, as opposed to another,
242                // unexpected cause.
243                mz_ore::process::exit_thread_safe(167);
244            }
245        } else {
246            // Reset burst budget if under limit.
247            self.burst_budget_remaining = self.config.burst_budget;
248        }
249
250        self.last_check = Instant::now();
251        Ok(())
252    }
253
254    /// Apply a new limiter config.
255    fn apply_config(&mut self, config: LimiterConfig) {
256        if config == self.config {
257            return; // no-op config change
258        }
259
260        info!(?config, "applying memory limiter config");
261        self.config = config;
262        self.burst_budget_remaining = config.burst_budget;
263
264        self.metrics
265            .memory_limit
266            .set(u64::cast_from(config.memory_limit));
267    }
268}
269
270struct LimiterMetrics {
271    duration: Histogram,
272    memory_limit: UIntGauge,
273    memory_usage: UIntGauge,
274    vm_rss: UIntGauge,
275    vm_swap: UIntGauge,
276    burst_budget: UIntGauge,
277}
278
279impl LimiterMetrics {
280    fn new(registry: &MetricsRegistry) -> Self {
281        Self {
282            duration: registry.register(metric!(
283                name: "mz_memory_limiter_duration_seconds",
284                help: "A histogram of the time it took to run the memory limiter.",
285                buckets: mz_ore::stats::histogram_seconds_buckets(0.000_500, 32.),
286            )),
287            memory_limit: registry.register(metric!(
288                name: "mz_memory_limiter_memory_limit_bytes",
289                help: "The configured memory limit.",
290            )),
291            memory_usage: registry.register(metric!(
292                name: "mz_memory_limiter_memory_usage_bytes",
293                help: "The current memory usage.",
294            )),
295            vm_rss: registry.register(metric!(
296                name: "mz_memory_limiter_vm_rss_bytes",
297                help: "The current VmRSS metric.",
298            )),
299            vm_swap: registry.register(metric!(
300                name: "mz_memory_limiter_vm_swap_bytes",
301                help: "The current VmSwap metric.",
302            )),
303            burst_budget: registry.register(metric!(
304                name: "mz_memory_limiter_burst_budget_byteseconds",
305                help: "The remaining memory burst budget.",
306            )),
307        }
308    }
309}
310
311/// Helper for reading and parsing `/proc/self/status` on Linux.
312pub struct ProcStatus {
313    /// Resident Set Size (RSS) in bytes.
314    pub vm_rss: usize,
315    /// Swap memory in bytes.
316    pub vm_swap: usize,
317}
318
319impl ProcStatus {
320    /// Populate a new `ProcStatus` with information in /proc/self/status.
321    pub fn from_proc() -> anyhow::Result<Self> {
322        let contents = std::fs::read_to_string("/proc/self/status")?;
323        let mut vm_rss = 0;
324        let mut vm_swap = 0;
325
326        for line in contents.lines() {
327            if line.starts_with("VmRSS:") {
328                vm_rss = line
329                    .split_whitespace()
330                    .nth(1)
331                    .ok_or_else(|| anyhow::anyhow!("failed to parse VmRSS"))?
332                    .parse::<usize>()
333                    .context("failed to parse VmRSS")?
334                    * 1024
335            } else if line.starts_with("VmSwap:") {
336                vm_swap = line
337                    .split_whitespace()
338                    .nth(1)
339                    .ok_or_else(|| anyhow::anyhow!("failed to parse VmSwap"))?
340                    .parse::<usize>()
341                    .context("failed to parse VmSwap")?
342                    * 1024;
343            }
344        }
345
346        Ok(Self { vm_rss, vm_swap })
347    }
348}
mz_compute/memory_limiter.rs

mz_compute/
memory_limiter.rs
