kumo_server_memory/
lib.rs

1//! This module contains logic to reason about memory usage,
2//! implement memory limits, and some helpers to attempt to
3//! release cached memory back to the system when memory
4//! is low.
5
6use crate::tracking::TrackingAllocator;
7use anyhow::Context;
8#[cfg(target_os = "linux")]
9use cgroups_rs::cgroup::{get_cgroups_relative_paths, Cgroup, UNIFIED_MOUNTPOINT};
10#[cfg(target_os = "linux")]
11use cgroups_rs::hierarchies::{V1, V2};
12#[cfg(target_os = "linux")]
13use cgroups_rs::memory::MemController;
14#[cfg(target_os = "linux")]
15use cgroups_rs::{Hierarchy, MaxValue};
16use kumo_prometheus::declare_metric;
17use nix::sys::resource::{rlim_t, RLIM_INFINITY};
18use nix::unistd::{sysconf, SysconfVar};
19#[cfg(any(target_os = "linux", test))]
20use std::collections::HashMap;
21use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
22use std::sync::{LazyLock, Mutex};
23use std::time::Duration;
24use tikv_jemallocator::Jemalloc;
25use tokio::sync::watch::Receiver;
26
27pub mod tracking;
28
29pub use tracking::{set_tracking_callstacks, tracking_stats};
30
31#[global_allocator]
32static GLOBAL: TrackingAllocator<Jemalloc> = TrackingAllocator::new(Jemalloc);
33
34declare_metric! {
35/// How many times the low memory threshold was exceeded.
36static LOW_COUNT: IntCounter("memory_low_count");
37}
38
39declare_metric! {
40/// how many times the soft memory limit was exceeded
41static OVER_LIMIT_COUNT: IntCounter("memory_over_limit_count");
42}
43
44declare_metric! {
45/// number of bytes of used memory that drives memory limit decisions.
46/// When sourced from a cgroup this is the working set
47/// (memory.current - inactive_file, floored by anonymous memory);
48/// otherwise it is the Resident Set Size from /proc/self/statm.
49static MEM_USAGE: Gauge("memory_usage");
50}
51
52declare_metric! {
53/// soft memory limit measured in bytes
54static MEM_LIMIT: Gauge("memory_limit");
55}
56
57declare_metric! {
58/// number of bytes of used memory (allocated by Rust)
59static MEM_COUNTED: Gauge("memory_usage_rust");
60}
61
62declare_metric! {
63/// low memory threshold measured in bytes
64static LOW_MEM_THRESH: Gauge("memory_low_thresh");
65}
66
67static SUBSCRIBER: LazyLock<Mutex<Option<Receiver<()>>>> = LazyLock::new(|| Mutex::new(None));
68
69static OVER_LIMIT: AtomicBool = AtomicBool::new(false);
70static LOW_MEM: AtomicBool = AtomicBool::new(false);
71
72/// Set true once we have warned about an unreadable cgroup memory.stat,
73/// so the fail-safe path warns at most once.
74#[cfg(target_os = "linux")]
75static WARNED_EMPTY_MEMORY_STAT: AtomicBool = AtomicBool::new(false);
76
77// Default this to a reasonable non-zero value, as it is possible
78// in the test harness on the CI system to inject concurrent with
79// early startup.
80// Having this return 0 and propagate back as a 421 makes it harder
81// to to write tests that care precisely about a response if they
82// have to deal with this small window on startup.
83static HEAD_ROOM: AtomicUsize = AtomicUsize::new(u32::MAX as usize);
84
85/// Represents the current memory usage of this process.
86///
87/// `bytes` is the value used for all memory-limit decisions. When sourced from
88/// a cgroup it is the working set (`current - inactive_file`, floored by
89/// anonymous memory); otherwise it is the Resident Set Size.
90#[derive(Debug, Clone, Copy)]
91pub struct MemoryUsage {
92    pub bytes: u64,
93}
94
95impl std::fmt::Display for MemoryUsage {
96    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
97        write!(fmt, "{}", human(self.bytes))
98    }
99}
100
101impl MemoryUsage {
102    pub fn get() -> anyhow::Result<Self> {
103        #[cfg(target_os = "linux")]
104        {
105            if let Ok(v2) = Self::get_cgroup(true) {
106                return Ok(v2);
107            }
108            if let Ok(v1) = Self::get_cgroup(false) {
109                return Ok(v1);
110            }
111        }
112        Self::get_linux_statm()
113    }
114
115    #[cfg(target_os = "linux")]
116    fn get_cgroup(v2: bool) -> anyhow::Result<Self> {
117        let cgroup = get_my_cgroup(v2)?;
118        let mem: &MemController = cgroup
119            .controller_of()
120            .ok_or_else(|| anyhow::anyhow!("no memory controller?"))?;
121        let stat = mem.memory_stat();
122        let current = stat.usage_in_bytes;
123
124        match parse_working_set_inputs(v2, &stat.stat.raw) {
125            Some((inactive_file, anon)) => Ok(Self {
126                bytes: working_set_usage(current, inactive_file, anon),
127            }),
128            None => {
129                // memory.stat was empty/unreadable: fall back to the raw
130                // usage rather than erroring, and warn at most once.
131                if !WARNED_EMPTY_MEMORY_STAT.swap(true, Ordering::Relaxed) {
132                    tracing::warn!(
133                        "cgroup memory.stat was empty or unreadable; \
134                         reporting raw memory.current ({}) as usage",
135                        human(current)
136                    );
137                }
138                Ok(Self { bytes: current })
139            }
140        }
141    }
142
143    pub fn get_linux_statm() -> anyhow::Result<Self> {
144        let data = std::fs::read_to_string("/proc/self/statm")?;
145        let fields: Vec<&str> = data.split(' ').collect();
146        let rss: u64 = fields[1].parse()?;
147        Ok(Self {
148            bytes: rss * sysconf(SysconfVar::PAGE_SIZE)?.unwrap_or(4 * 1024) as u64,
149        })
150    }
151}
152
153/// Computes the working set from cgroup memory statistics.
154///
155/// `current - inactive_file` removes cold, kernel-reclaimable page cache;
156/// the `max(anon)` floor ensures a fast anonymous burst between samples
157/// cannot hide behind a transiently low working set. Saturating subtraction
158/// guards against `inactive_file` exceeding `current`.
159#[cfg(any(target_os = "linux", test))]
160fn working_set_usage(current: u64, inactive_file: u64, anon: u64) -> u64 {
161    current.saturating_sub(inactive_file).max(anon)
162}
163
164/// Extracts `(inactive_file, anon)` from a parsed `memory.stat` for the
165/// working-set computation, or `None` when `memory.stat` could not be read
166/// (the crate returns an empty `raw` map on failure).
167///
168/// cgroup v2 uses the `inactive_file` / `anon` keys. cgroup v1 uses the
169/// hierarchical `total_inactive_file` key, and `anon` is
170/// `total_rss + total_rss_huge`; the hierarchical (`total_*`) values are used
171/// so the anon floor is consistent with the hierarchical `memory.usage_in_bytes`
172/// it is compared against.
173#[cfg(any(target_os = "linux", test))]
174fn parse_working_set_inputs(v2: bool, raw: &HashMap<String, u64>) -> Option<(u64, u64)> {
175    if raw.is_empty() {
176        return None;
177    }
178    let get = |key: &str| raw.get(key).copied().unwrap_or(0);
179    if v2 {
180        Some((get("inactive_file"), get("anon")))
181    } else {
182        let anon = get("total_rss").saturating_add(get("total_rss_huge"));
183        Some((get("total_inactive_file"), anon))
184    }
185}
186
187fn human(n: u64) -> String {
188    humansize::format_size(n, humansize::DECIMAL)
189}
190
191/// Represents a constraint on memory usage
192#[derive(Debug, Clone, Copy)]
193pub struct MemoryLimits {
194    pub soft_limit: Option<u64>,
195    pub hard_limit: Option<u64>,
196}
197
198impl std::fmt::Display for MemoryLimits {
199    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
200        let soft = self.soft_limit.map(human);
201        let hard = self.hard_limit.map(human);
202        write!(fmt, "soft={soft:?}, hard={hard:?}")
203    }
204}
205
206impl MemoryLimits {
207    pub fn min(self, other: Self) -> Self {
208        Self {
209            soft_limit: min_opt_limit(self.soft_limit, other.soft_limit),
210            hard_limit: min_opt_limit(self.hard_limit, other.hard_limit),
211        }
212    }
213
214    pub fn is_unlimited(&self) -> bool {
215        self.soft_limit.is_none() && self.hard_limit.is_none()
216    }
217}
218
219fn rlim_to_opt(rlim: rlim_t) -> Option<u64> {
220    if rlim == RLIM_INFINITY {
221        None
222    } else {
223        Some(rlim)
224    }
225}
226
227#[cfg(target_os = "linux")]
228fn max_value_to_opt(value: Option<MaxValue>) -> anyhow::Result<Option<u64>> {
229    Ok(match value {
230        None | Some(MaxValue::Max) => None,
231        Some(MaxValue::Value(n)) if n >= 0 => Some(n as u64),
232        Some(MaxValue::Value(n)) => anyhow::bail!("unexpected negative limit {n}"),
233    })
234}
235
236fn min_opt_limit(a: Option<u64>, b: Option<u64>) -> Option<u64> {
237    match (a, b) {
238        (Some(a), Some(b)) => Some(a.min(b)),
239        (Some(a), None) | (None, Some(a)) => Some(a),
240        (None, None) => None,
241    }
242}
243
244impl MemoryLimits {
245    pub fn get_rlimits() -> anyhow::Result<Self> {
246        #[cfg(not(target_os = "macos"))]
247        let (rss_soft, rss_hard) =
248            nix::sys::resource::getrlimit(nix::sys::resource::Resource::RLIMIT_RSS)?;
249        #[cfg(target_os = "macos")]
250        let (rss_soft, rss_hard) = (RLIM_INFINITY, RLIM_INFINITY);
251
252        let soft_limit = rlim_to_opt(rss_soft);
253        let hard_limit = rlim_to_opt(rss_hard);
254
255        Ok(Self {
256            soft_limit,
257            hard_limit,
258        })
259    }
260
261    #[cfg(target_os = "linux")]
262    fn get_any_cgroup() -> anyhow::Result<Self> {
263        if let Ok(cg) = Self::get_cgroup(true) {
264            return Ok(cg);
265        }
266        Self::get_cgroup(false)
267    }
268
269    #[cfg(target_os = "linux")]
270    pub fn get_cgroup(v2: bool) -> anyhow::Result<Self> {
271        let cgroup = get_my_cgroup(v2)?;
272        let mem: &MemController = cgroup
273            .controller_of()
274            .ok_or_else(|| anyhow::anyhow!("no memory controller?"))?;
275
276        let limits = mem.get_mem()?;
277        Ok(Self {
278            soft_limit: max_value_to_opt(limits.high)?,
279            hard_limit: max_value_to_opt(limits.max)?,
280        })
281    }
282}
283
284/// Returns the amount of physical memory available to the system.
285/// This is linux specific.
286#[cfg(target_os = "linux")]
287fn get_physical_memory() -> anyhow::Result<u64> {
288    let data = std::fs::read_to_string("/proc/meminfo")?;
289    for line in data.lines() {
290        if line.starts_with("MemTotal:") {
291            let mut iter = line.rsplit(' ');
292            let unit = iter
293                .next()
294                .ok_or_else(|| anyhow::anyhow!("expected unit"))?;
295            if unit != "kB" {
296                anyhow::bail!("unsupported /proc/meminfo unit {unit}");
297            }
298            let value = iter
299                .next()
300                .ok_or_else(|| anyhow::anyhow!("expected value"))?;
301            let value: u64 = value.parse()?;
302
303            return Ok(value * 1024);
304        }
305    }
306    anyhow::bail!("MemTotal not found in /proc/meminfo");
307}
308
309/// Retrieves the current usage and limits.
310/// This is a bit of a murky area as, on Linux, the cgroup reported usage
311/// appears to be nonsensical when no limits are configured.
312/// So we first obtain the limits from cgroups, and if they are set,
313/// we return the usage from cgroups along with it,
314/// otherwise we get the ulimit limits and look at the more general
315/// usage numbers to go with it.
316///
317/// If no limits are explicitly configured, we'll assume a hard limit
318/// equal to the physical ram on the system, and a soft limit of 75%
319/// of whatever we've determined the hard limit to be.
320#[cfg(target_os = "linux")]
321fn get_usage_and_limit_impl() -> anyhow::Result<(MemoryUsage, MemoryLimits)> {
322    let mut limit = MemoryLimits::get_rlimits()?;
323    let mut usage = MemoryUsage::get_linux_statm()?;
324
325    if let Ok(cg_lim) = MemoryLimits::get_any_cgroup() {
326        if !cg_lim.is_unlimited() {
327            limit = limit.min(cg_lim);
328            usage = MemoryUsage::get()?;
329        }
330    }
331
332    if limit.hard_limit.is_none() {
333        let phys = get_physical_memory()?;
334        limit.hard_limit.replace(phys);
335    }
336    if limit.soft_limit.is_none() {
337        limit.soft_limit = limit.hard_limit.map(|lim| lim * 3 / 4);
338    }
339
340    Ok((usage, limit))
341}
342
343#[cfg(not(target_os = "linux"))]
344fn get_usage_and_limit_impl() -> anyhow::Result<(MemoryUsage, MemoryLimits)> {
345    Ok((
346        MemoryUsage { bytes: 0 },
347        MemoryLimits {
348            soft_limit: None,
349            hard_limit: None,
350        },
351    ))
352}
353
354static USER_HARD_LIMIT: AtomicUsize = AtomicUsize::new(0);
355static USER_SOFT_LIMIT: AtomicUsize = AtomicUsize::new(0);
356static USER_LOW_THRESH: AtomicUsize = AtomicUsize::new(0);
357
358pub fn set_hard_limit(limit: usize) {
359    USER_HARD_LIMIT.store(limit, Ordering::Relaxed);
360}
361
362pub fn set_soft_limit(limit: usize) {
363    USER_SOFT_LIMIT.store(limit, Ordering::Relaxed);
364}
365
366pub fn set_low_memory_thresh(limit: usize) {
367    USER_LOW_THRESH.store(limit, Ordering::Relaxed);
368}
369
370pub fn get_hard_limit() -> Option<u64> {
371    let user_limit = USER_HARD_LIMIT.load(Ordering::Relaxed);
372    if user_limit > 0 {
373        return Some(user_limit as u64);
374    }
375
376    let result = get_usage_and_limit()
377        .ok()
378        .and_then(|(_, limits)| limits.hard_limit);
379    result
380}
381
382pub fn get_soft_limit() -> Option<u64> {
383    let user_limit = USER_SOFT_LIMIT.load(Ordering::Relaxed);
384    if user_limit > 0 {
385        return Some(user_limit as u64);
386    }
387
388    get_usage_and_limit()
389        .ok()
390        .and_then(|(_, limits)| limits.soft_limit)
391}
392
393pub fn get_low_memory_thresh() -> Option<u64> {
394    let user_thresh = USER_LOW_THRESH.load(Ordering::Relaxed);
395    if user_thresh > 0 {
396        return Some(user_thresh as u64);
397    }
398
399    get_usage_and_limit()
400        .ok()
401        .and_then(|(_, limits)| limits.soft_limit)
402        .map(|limit| limit * 8 / 10)
403}
404
405pub fn get_usage_and_limit() -> anyhow::Result<(MemoryUsage, MemoryLimits)> {
406    let (usage, mut limit) = get_usage_and_limit_impl()?;
407
408    if let Ok(value) = std::env::var("KUMOD_MEMORY_HARD_LIMIT") {
409        limit.hard_limit.replace(
410            value
411                .parse()
412                .context("failed to parse KUMOD_MEMORY_HARD_LIMIT env var")?,
413        );
414    }
415    if let Ok(value) = std::env::var("KUMOD_MEMORY_SOFT_LIMIT") {
416        limit.soft_limit.replace(
417            value
418                .parse()
419                .context("failed to parse KUMOD_MEMORY_SOFT_LIMIT env var")?,
420        );
421    }
422
423    let hard = USER_HARD_LIMIT.load(Ordering::Relaxed);
424    if hard > 0 {
425        limit.hard_limit.replace(hard as u64);
426    }
427    let soft = USER_SOFT_LIMIT.load(Ordering::Relaxed);
428    if soft > 0 {
429        limit.soft_limit.replace(soft as u64);
430    }
431
432    Ok((usage, limit))
433}
434
435/// To be called when a thread goes idle; it will flush cached
436/// memory out of the thread local cache to be returned/reused
437/// elsewhere in the system
438pub fn purge_thread_cache() {
439    unsafe {
440        tikv_jemalloc_sys::mallctl(
441            b"thread.tcache.flush\0".as_ptr() as *const _,
442            std::ptr::null_mut(),
443            std::ptr::null_mut(),
444            std::ptr::null_mut(),
445            0,
446        );
447    }
448}
449
450/// To be called when used memory is high: will aggressively
451/// flush and release cached memory
452fn purge_all_arenas() {
453    unsafe {
454        // 4096 is MALLCTL_ARENAS_ALL, which is a magic value
455        // that instructs jemalloc to purge all arenas
456        tikv_jemalloc_sys::mallctl(
457            b"arena.4096.purge\0".as_ptr() as *const _,
458            std::ptr::null_mut(),
459            std::ptr::null_mut(),
460            std::ptr::null_mut(),
461            0,
462        );
463    }
464}
465
466/// If `MALLOC_CONF='prof:true,prof_prefix:jeprof.out'` is set in the
467/// environment, calling this will generate a heap profile in the
468/// current directory
469fn dump_heap_profile() {
470    unsafe {
471        tikv_jemalloc_sys::mallctl(
472            b"prof.dump\0".as_ptr() as *const _,
473            std::ptr::null_mut(),
474            std::ptr::null_mut(),
475            std::ptr::null_mut(),
476            0,
477        );
478    }
479}
480
481/// The memory thread continuously examines memory usage and limits
482/// and maintains global counters to track the memory state
483fn memory_thread() {
484    let mut is_ok = true;
485    let mut is_low = false;
486
487    let (tx, rx) = tokio::sync::watch::channel(());
488    SUBSCRIBER.lock().unwrap().replace(rx);
489
490    loop {
491        MEM_COUNTED.set(crate::tracking::counted_usage() as f64);
492
493        match get_usage_and_limit() {
494            Ok((
495                MemoryUsage { bytes: usage },
496                MemoryLimits {
497                    soft_limit: Some(limit),
498                    hard_limit: _,
499                },
500            )) => {
501                let was_ok = is_ok;
502                is_ok = usage < limit;
503                OVER_LIMIT.store(is_ok, Ordering::SeqCst);
504                HEAD_ROOM.store(limit.saturating_sub(usage) as usize, Ordering::SeqCst);
505                MEM_USAGE.set(usage as f64);
506                MEM_LIMIT.set(limit as f64);
507
508                let mut low_thresh = USER_LOW_THRESH.load(Ordering::Relaxed) as u64;
509                if low_thresh == 0 {
510                    low_thresh = limit * 8 / 10;
511                }
512                LOW_MEM_THRESH.set(low_thresh as f64);
513
514                let was_low = is_low;
515                is_low = usage > low_thresh;
516                LOW_MEM.store(is_low, Ordering::SeqCst);
517
518                if !was_low && is_low {
519                    // Transition to low memory
520                    LOW_COUNT.inc();
521                }
522
523                if !is_ok && was_ok {
524                    // Transition from OK -> !OK
525                    dump_heap_profile();
526                    OVER_LIMIT_COUNT.inc();
527                    tracing::error!(
528                        "memory usage {} exceeds limit {}",
529                        human(usage),
530                        human(limit)
531                    );
532                    tx.send(()).ok();
533                    purge_all_arenas();
534                } else if !was_ok && is_ok {
535                    // Transition from !OK -> OK
536                    dump_heap_profile();
537                    tracing::error!(
538                        "memory usage {} is back within limit {}",
539                        human(usage),
540                        human(limit)
541                    );
542                    tx.send(()).ok();
543                } else {
544                    if !is_ok {
545                        purge_all_arenas();
546                    }
547                    tracing::debug!("memory usage {}, limit {}", human(usage), human(limit));
548                }
549            }
550            Ok((
551                MemoryUsage { bytes: 0 },
552                MemoryLimits {
553                    soft_limit: None,
554                    hard_limit: None,
555                },
556            )) => {
557                // We don't know anything about the memory usage on this
558                // system, just pretend everything is fine
559                HEAD_ROOM.store(1024, Ordering::SeqCst);
560            }
561            Ok(_) => {}
562            Err(err) => tracing::error!("unable to query memory info: {err:#}"),
563        }
564
565        std::thread::sleep(Duration::from_secs(3));
566    }
567}
568
569/// Returns the amount of headroom; the number of bytes that can
570/// be allocated before we hit the soft limit
571pub fn get_headroom() -> usize {
572    HEAD_ROOM.load(Ordering::SeqCst)
573}
574
575/// Returns true when we are within 10% if the soft limit
576pub fn low_memory() -> bool {
577    LOW_MEM.load(Ordering::SeqCst)
578}
579
580/// Indicates the overall memory status
581pub fn memory_status() -> MemoryStatus {
582    if get_headroom() == 0 {
583        MemoryStatus::NoMemory
584    } else if low_memory() {
585        MemoryStatus::LowMemory
586    } else {
587        MemoryStatus::Ok
588    }
589}
590
591#[derive(Copy, Clone, Debug, Eq, PartialEq)]
592pub enum MemoryStatus {
593    Ok,
594    LowMemory,
595    NoMemory,
596}
597
598/// Returns a receiver that will notify when memory status
599/// changes from OK -> !OK or vice versa.
600pub fn subscribe_to_memory_status_changes() -> Option<Receiver<()>> {
601    SUBSCRIBER.lock().unwrap().clone()
602}
603
604pub async fn subscribe_to_memory_status_changes_async() -> Receiver<()> {
605    loop {
606        if let Some(rx) = subscribe_to_memory_status_changes() {
607            return rx;
608        }
609        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
610    }
611}
612
613/// Initialize the memory thread to monitor memory usage/limits
614pub fn setup_memory_limit() -> anyhow::Result<()> {
615    let (usage, limit) = get_usage_and_limit()?;
616    tracing::debug!("usage: {usage:?}");
617    tracing::info!("using limits: {limit}");
618
619    std::thread::Builder::new()
620        .name("memory-monitor".to_string())
621        .spawn(memory_thread)?;
622
623    Ok(())
624}
625
626/// Returns a Cgroup for the current process.
627/// Can return either a v2 or a v1 cgroup.
628#[cfg(target_os = "linux")]
629fn get_my_cgroup(v2: bool) -> anyhow::Result<Cgroup> {
630    let paths = get_cgroups_relative_paths()?;
631    let h: Box<dyn Hierarchy> = if v2 {
632        Box::new(V2::new())
633    } else {
634        Box::new(V1::new())
635    };
636
637    let path = paths
638        .get("")
639        .ok_or_else(|| anyhow::anyhow!("couldn't resolve path"))?;
640
641    let cgroup = Cgroup::load(h, format!("{}/{}", UNIFIED_MOUNTPOINT, path));
642    Ok(cgroup)
643}
644
645#[derive(Copy, Clone)]
646pub struct NumBytes(pub usize);
647
648impl std::fmt::Debug for NumBytes {
649    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
650        write!(
651            fmt,
652            "{} ({})",
653            self.0,
654            humansize::format_size(self.0, humansize::DECIMAL)
655        )
656    }
657}
658
659impl From<usize> for NumBytes {
660    fn from(n: usize) -> Self {
661        Self(n)
662    }
663}
664
665impl From<u64> for NumBytes {
666    fn from(n: u64) -> Self {
667        Self(n as usize)
668    }
669}
670
671#[derive(Copy, Clone)]
672pub struct Number(pub usize);
673
674impl std::fmt::Debug for Number {
675    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
676        use num_format::{Locale, ToFormattedString};
677        write!(
678            fmt,
679            "{} ({})",
680            self.0,
681            self.0.to_formatted_string(&Locale::en)
682        )
683    }
684}
685
686impl From<usize> for Number {
687    fn from(n: usize) -> Self {
688        Self(n)
689    }
690}
691
692impl From<u64> for Number {
693    fn from(n: u64) -> Self {
694        Self(n as usize)
695    }
696}
697
698#[derive(Debug)]
699pub struct JemallocStats {
700    /// stats.allocated` - Total number of bytes allocated by the application.
701    pub allocated: NumBytes,
702
703    /// `stats.active`
704    /// Total number of bytes in active pages allocated by the application. This is a multiple of
705    /// the page size, and greater than or equal to stats.allocated. This does not include
706    /// `stats.arenas.<i>.pdirty`, `stats.arenas.<i>.pmuzzy`, nor pages entirely devoted to allocator
707    /// metadata.
708    pub active: NumBytes,
709
710    /// stats.metadata (size_t) r- [--enable-stats]
711    /// Total number of bytes dedicated to metadata, which comprise base allocations used for bootstrap-sensitive allocator metadata structures (see `stats.arenas.<i>.base`) and internal allocations (see `stats.arenas.<i>.internal`). Transparent huge page (enabled with opt.metadata_thp) usage is not considered.
712    pub metadata: NumBytes,
713
714    /// stats.resident (size_t) r- [--enable-stats]
715    /// Maximum number of bytes in physically resident data pages mapped by the allocator, comprising all pages dedicated to allocator metadata, pages backing active allocations, and unused dirty pages. This is a maximum rather than precise because pages may not actually be physically resident if they correspond to demand-zeroed virtual memory that has not yet been touched. This is a multiple of the page size, and is larger than stats.active.
716    pub resident: NumBytes,
717
718    /// stats.mapped (size_t) r- [--enable-stats]
719    /// Total number of bytes in active extents mapped by the allocator. This is larger than stats.active. This does not include inactive extents, even those that contain unused dirty pages, which means that there is no strict ordering between this and stats.resident.
720    pub mapped: NumBytes,
721
722    /// stats.retained (size_t) r- [--enable-stats]
723    /// Total number of bytes in virtual memory mappings that were retained rather than being returned to the operating system via e.g. munmap(2) or similar. Retained virtual memory is typically untouched, decommitted, or purged, so it has no strongly associated physical memory (see extent hooks for details). Retained memory is excluded from mapped memory statistics, e.g. stats.mapped.
724    pub retained: NumBytes,
725}
726
727impl JemallocStats {
728    pub fn collect() -> Self {
729        use tikv_jemalloc_ctl::{epoch, stats};
730
731        // jemalloc stats are cached and don't fully report current
732        // values until an epoch is advanced, so let's explicitly
733        // do that here.
734        epoch::advance().ok();
735
736        Self {
737            allocated: stats::allocated::read().unwrap_or(0).into(),
738            active: stats::active::read().unwrap_or(0).into(),
739            metadata: stats::metadata::read().unwrap_or(0).into(),
740            resident: stats::resident::read().unwrap_or(0).into(),
741            mapped: stats::mapped::read().unwrap_or(0).into(),
742            retained: stats::retained::read().unwrap_or(0).into(),
743        }
744    }
745}
746
747#[cfg(test)]
748mod tests {
749    use super::*;
750
751    const GB: u64 = 1024 * 1024 * 1024;
752    const MB: u64 = 1024 * 1024;
753
754    fn raw(pairs: &[(&str, u64)]) -> HashMap<String, u64> {
755        pairs
756            .iter()
757            .map(|(key, value)| ((*key).to_string(), *value))
758            .collect()
759    }
760
761    #[test]
762    fn working_set_normal() {
763        // The motivating pod: 51GB current, 46GB cold cache, 1.8GB real
764        // footprint. Working set should drop the cold cache.
765        let current = 51 * GB;
766        let inactive_file = 46 * GB;
767        let anon = 1_800 * MB;
768        assert_eq!(
769            working_set_usage(current, inactive_file, anon),
770            current - inactive_file
771        );
772    }
773
774    #[test]
775    fn anon_floor_wins() {
776        // A fast anon burst exceeds the working set; the floor must win so a
777        // transiently low working set can't hide it.
778        let current = 5 * GB;
779        let inactive_file = 4 * GB;
780        let anon = 4 * GB;
781        assert_eq!(working_set_usage(current, inactive_file, anon), 4 * GB);
782    }
783
784    #[test]
785    fn underflow_no_panic() {
786        // inactive_file >= current must saturate to the anon floor (0 here),
787        // never panic.
788        assert_eq!(working_set_usage(10 * GB, 10 * GB, 0), 0);
789        assert_eq!(working_set_usage(10 * GB, 12 * GB, 0), 0);
790    }
791
792    #[test]
793    fn parse_v2() {
794        // active_file / file are present but irrelevant; only inactive_file
795        // and anon feed the working set.
796        let raw = raw(&[
797            ("anon", 2 * GB),
798            ("inactive_file", 3 * GB),
799            ("active_file", GB),
800            ("file", 4 * GB),
801        ]);
802        assert_eq!(parse_working_set_inputs(true, &raw), Some((3 * GB, 2 * GB)));
803    }
804
805    #[test]
806    fn parse_v1_hierarchical() {
807        // v1 uses the hierarchical total_* keys, and anon is
808        // total_rss + total_rss_huge. The non-hierarchical rss must be ignored.
809        let raw = raw(&[
810            ("total_rss", 2 * GB),
811            ("total_rss_huge", 512 * MB),
812            ("total_inactive_file", 3 * GB),
813            ("rss", GB),
814        ]);
815        assert_eq!(
816            parse_working_set_inputs(false, &raw),
817            Some((3 * GB, 2 * GB + 512 * MB))
818        );
819    }
820
821    #[test]
822    fn empty_raw_is_none() {
823        // An empty raw map is how the crate signals an unreadable memory.stat;
824        // both versions must report None so the caller can fail safe.
825        let raw = HashMap::new();
826        assert_eq!(parse_working_set_inputs(true, &raw), None);
827        assert_eq!(parse_working_set_inputs(false, &raw), None);
828    }
829
830    #[test]
831    fn missing_keys_default_zero() {
832        // A non-empty map missing our keys yields zeros (still Some), so the
833        // working set degrades to current rather than failing safe.
834        let raw = raw(&[("cache", GB)]);
835        assert_eq!(parse_working_set_inputs(true, &raw), Some((0, 0)));
836    }
837}
kumo_server_memory/lib.rs

kumo_server_memory/
lib.rs
