process_memory/

linux.rs

use std::{
    fs::{self, File},
    io::{self, Read},
    mem::MaybeUninit,
};

const SMAPS_ROLLUP_PATH: &str = "/proc/self/smaps_rollup";
const SMAPS_PATH: &str = "/proc/self/smaps";
const STATM: &str = "/proc/self/statm";
const RSS_LINE_PREFIX: &[u8] = b"Rss: ";

enum StatSource {
    SmapsRollup(Scanner<File>),
    Smaps(Scanner<File>),
    Statm(Option<usize>),
}

/// A memory usage querier.
pub struct Querier {
    source: StatSource,
}

impl Querier {
    /// Gets the resident set size of this process, in bytes.
    ///
    /// If the resident set size can't be determined, `None` is returned. This could be for a number of underlying
    /// reasons, but should generally be considered an incredibly rare/unlikely event.
    pub fn resident_set_size(&mut self) -> Option<usize> {
        match &mut self.source {
            StatSource::SmapsRollup(scanner) => {
                // As smaps_rollup is a pre-aggregated version of smaps, there's only one "Rss:" line that we need to find, so use
                // the same scanner-based approach as we do for smaps, but just take the first matching line we find.
                scanner.reset_with_path(SMAPS_ROLLUP_PATH).ok()?;
                first_rss_value(scanner)
            }
            StatSource::Smaps(scanner) => {
                // Scan all lines in smaps, looking for lines that start with "Rss:". Each of these lines will contain the resident
                // set size of a particular memory mapping. We simply need to find all of these lines and aggregate their value to
                // get the RSS for the process.
                scanner.reset_with_path(SMAPS_PATH).ok()?;
                sum_rss_values(scanner)
            }
            StatSource::Statm(maybe_page_size) => {
                let page_size = maybe_page_size.as_ref().copied()?;

                // Unlike smaps/smaps_rollup, statm is a drastically simpler format that is written as a single line with
                // space-delimited fields. Since we have no lines to scan, it's much simpler to just read the entire file into a
                // stack-allocated buffer.
                //
                // With seven integer fields, we can napkin math this to wanting to hold 20 bytes per field, plus the separators and
                // newline, which is 153... but power-of-two numbers somehow feel better, so we'll go to 256.
                let mut buf = [0; 256];
                let mut file = File::open(STATM).ok()?;
                let n = file.read(&mut buf).ok()?;
                if n == 0 || n == buf.len() {
                    // If we read no bytes, or filled the entire buffer, something is very wrong.
                    return None;
                }

                parse_statm_rss_bytes(&buf, page_size)
            }
        }
    }
}

/// Returns the first `Rss:` value (converted to bytes) found by the scanner, or `None` if there is none.
///
/// This is the `smaps_rollup` strategy: the file is pre-aggregated, so the first matching line is the total RSS.
fn first_rss_value<T: Read>(scanner: &mut Scanner<T>) -> Option<usize> {
    while let Ok(Some(raw_rss_line)) = scanner.next_matching_line(RSS_LINE_PREFIX) {
        let raw_rss_value = skip_to_line_value(raw_rss_line)?;
        if let Some(rss_bytes) = parse_kb_value_as_bytes(raw_rss_value) {
            return Some(rss_bytes);
        }
    }

    None
}

/// Returns the sum of every `Rss:` value (converted to bytes) found by the scanner, or `None` if the total is zero.
///
/// This is the `smaps` strategy: each memory mapping contributes its own `Rss:` line, and the process RSS is their sum.
fn sum_rss_values<T: Read>(scanner: &mut Scanner<T>) -> Option<usize> {
    let mut total_rss_bytes = 0;
    while let Ok(Some(raw_rss_line)) = scanner.next_matching_line(RSS_LINE_PREFIX) {
        let raw_rss_value = skip_to_line_value(raw_rss_line)?;
        if let Some(rss_bytes) = parse_kb_value_as_bytes(raw_rss_value) {
            total_rss_bytes += rss_bytes;
        }
    }

    if total_rss_bytes > 0 {
        Some(total_rss_bytes)
    } else {
        None
    }
}

/// Parses the resident set size (in bytes) out of the contents of a `statm` file.
///
/// The resident set size is the second whitespace-delimited field, expressed in pages, so it is multiplied by the page
/// size to get a byte count. `None` is returned if the field is missing or non-numeric.
fn parse_statm_rss_bytes(raw: &[u8], page_size: usize) -> Option<usize> {
    // Resident set size is the second field, so we need to skip to it.
    let raw_rss_field = raw.split(|b| *b == b' ').nth(1)?;

    // We need to parse the field as an integer, and then multiply it by the page size to get the value in bytes.
    let rss_pages = simdutf8::basic::from_utf8(raw_rss_field).ok()?.parse::<usize>().ok()?;
    Some(rss_pages * page_size)
}

impl Default for Querier {
    fn default() -> Self {
        Self {
            source: determine_stat_source(),
        }
    }
}

fn determine_stat_source() -> StatSource {
    select_stat_source(
        fs::metadata(SMAPS_ROLLUP_PATH).is_ok(),
        fs::metadata(SMAPS_PATH).is_ok(),
        page_size(),
    )
}

/// Selects the RSS data source based on which procfs files are available.
///
/// The preference order (documented at the crate level) is `smaps_rollup`, then `smaps`, then `statm`. Splitting the
/// availability checks from the selection logic keeps the fallback order testable without touching the real filesystem.
fn select_stat_source(has_smaps_rollup: bool, has_smaps: bool, page_size: Option<usize>) -> StatSource {
    if has_smaps_rollup {
        StatSource::SmapsRollup(Scanner::new())
    } else if has_smaps {
        StatSource::Smaps(Scanner::new())
    } else {
        StatSource::Statm(page_size)
    }
}

fn page_size() -> Option<usize> {
    let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
    if page_size <= 0 {
        None
    } else {
        Some(page_size as usize)
    }
}

fn parse_kb_value_as_bytes(raw_rss_value: &[u8]) -> Option<usize> {
    // The raw value here will be in the form of `XXXXXX kB`, so we want to find the first whitespace character, and
    // take everything before that.
    match raw_rss_value.iter().position(|&b| b == b' ') {
        Some(space_idx) => {
            let raw_value = &raw_rss_value[..space_idx];
            simdutf8::basic::from_utf8(raw_value)
                .ok()?
                .parse::<usize>()
                .ok()
                .map(|value| value * 1024)
        }
        None => None,
    }
}

struct Scanner<T> {
    io: Option<T>,
    eof: bool,
    buf: Vec<u8>,
    pending_consume: Option<usize>,
}

impl<T> Scanner<T>
where
    T: Read,
{
    fn new() -> Self {
        Self {
            io: None,
            eof: false,
            buf: Vec::with_capacity(8192),
            pending_consume: None,
        }
    }

    fn reset(&mut self, io: T) {
        self.buf.clear();
        self.eof = false;
        self.pending_consume = None;
        self.io = Some(io);
    }

    fn get_io_mut(&mut self) -> io::Result<&mut T> {
        match self.io.as_mut() {
            Some(io) => Ok(io),
            None => Err(io::Error::other("no file set in scanner")),
        }
    }

    fn fill_buf(&mut self) -> io::Result<()> {
        if self.eof {
            return Ok(());
        }

        // If our buffer isn't entirely filled, try and fill the remainder.
        if self.buf.len() < self.buf.capacity() {
            // If we have any spare capacity, try to read as many bytes as we can hold in it.
            //
            // SAFETY: There's no invalid bit patterns for `u8`.
            let read_buf = unsafe { &mut *(self.buf.spare_capacity_mut() as *mut [MaybeUninit<u8>] as *mut [u8]) };
            let n = self.get_io_mut()?.read(read_buf)?;
            if n == 0 {
                self.eof = true;
            }

            // SAFETY: We've just read `n` bytes into `buf`, based on the spare capacity, so incrementing our length by
            // `n` will only cover initialized bytes, and can't result in a length greater than the buffer capacity.
            unsafe {
                self.buf.set_len(self.buf.len() + n);
            }
        }

        Ok(())
    }

    fn next_matching_line(&mut self, prefix: &[u8]) -> io::Result<Option<&[u8]>> {
        loop {
            // We've reached EOF and have processed the entire file.
            if self.eof && self.buf.is_empty() {
                return Ok(None);
            }

            // If we have a pending consume, take that many bytes from the front of the buffer and shift the rest of the
            // data forward.
            if let Some(consume) = self.pending_consume {
                self.buf.drain(..consume);
                self.pending_consume = None;
            }

            // Ensure our buffer is as filled as it can be.
            self.fill_buf()?;

            // Get the entire buffer that we currently have, and see if it starts with our prefix.
            //
            // If it does, we then need to also find a newline character to know where to chop it off.
            if self.buf.starts_with(prefix) {
                let maybe_newline_idx = self.buf.iter().position(|&b| b == b'\n');
                if let Some(newline_idx) = maybe_newline_idx {
                    // Consume up to and including the newline character, but only hand back the bytes up to the newline.
                    self.pending_consume = Some(newline_idx + 1);

                    return Ok(Some(&self.buf[..newline_idx]));
                }
            } else {
                // Our buffer doesn't start with the prefix, so we need to find the next newline character and consume
                // up to that point to reset ourselves.
                //
                // We're essentially resetting ourselves to start at the next line in the file.
                let maybe_newline_idx = self.buf.iter().position(|&b| b == b'\n');
                if let Some(newline_idx) = maybe_newline_idx {
                    self.pending_consume = Some(newline_idx + 1);
                } else {
                    // We couldn't find a newline character, so we need to clear the entire buffer and just keep reading.
                    self.buf.clear();
                }
            }
        }
    }
}

impl Scanner<File> {
    fn reset_with_path(&mut self, path: &str) -> io::Result<()> {
        let file = File::open(path)?;
        self.reset(file);

        Ok(())
    }
}

fn skip_to_line_value(raw_line: &[u8]) -> Option<&[u8]> {
    // We skip over all non-numeric characters and then return what's left.
    //
    // If the line doesn't contain any numeric characters, `None` is returned.
    raw_line
        .iter()
        .position(|b| b.is_ascii_digit())
        .map(|idx| &raw_line[idx..])
}

#[cfg(test)]
mod tests {
    use super::Querier;

    #[test]
    fn basic() {
        let mut querier = Querier::default();
        assert!(querier.resident_set_size().is_some());
    }

    #[test]
    fn skip_to_line_value() {
        let passing_lines = [
            "1234 kB".as_bytes(),
            "    1234 kB".as_bytes(),
            "\t1234 kB".as_bytes(),
            "Rss:1234 kB".as_bytes(),
            "Rss:   1234 kB".as_bytes(),
        ];
        for line in &passing_lines {
            assert_eq!(super::skip_to_line_value(line), Some("1234 kB".as_bytes()));
        }

        let failing_lines = [
            "Rss: ".as_bytes(),
            "Rss: \n".as_bytes(),
            "Rss:  kB".as_bytes(),
            "Rss: kB\n".as_bytes(),
        ];
        for line in &failing_lines {
            assert_eq!(super::skip_to_line_value(line), None);
        }
    }

    #[test]
    fn scanner_basic() {
        let prefix = "Rss: ".as_bytes();

        let mut scanner = super::Scanner::new();
        let mut buf = Vec::new();
        buf.extend_from_slice(b"Rss: 1234 kB\nRss: 5678 kB\nRss: 91011 kB\n");
        scanner.reset(buf.as_slice());

        assert_eq!(
            scanner.next_matching_line(prefix).unwrap(),
            Some(b"Rss: 1234 kB".as_ref())
        );
        assert_eq!(
            scanner.next_matching_line(prefix).unwrap(),
            Some(b"Rss: 5678 kB".as_ref())
        );
        assert_eq!(
            scanner.next_matching_line(prefix).unwrap(),
            Some(b"Rss: 91011 kB".as_ref())
        );
        assert_eq!(scanner.next_matching_line(prefix).unwrap(), None);
    }

    #[test]
    fn scanner_skip_non_matching_lines() {
        let prefix = "Rss: ".as_bytes();

        let mut scanner = super::Scanner::new();
        let mut buf = Vec::new();
        buf.extend_from_slice(b"Rss: 1234 kB\nPss: 5678 kB\nHugepages:    42069 kB\nRss: 91011 kB\n");
        scanner.reset(buf.as_slice());

        assert_eq!(
            scanner.next_matching_line(prefix).unwrap(),
            Some(b"Rss: 1234 kB".as_ref())
        );
        assert_eq!(
            scanner.next_matching_line(prefix).unwrap(),
            Some(b"Rss: 91011 kB".as_ref())
        );
        assert_eq!(scanner.next_matching_line(prefix).unwrap(), None);
    }

    #[test]
    fn scanner_skips_lines_larger_than_buffer() {
        let prefix = "Rss: ".as_bytes();

        let mut scanner = super::Scanner::new();

        // We construct a non-matching line that's longer than our internal buffer (8192 bytes) to ensure that we can
        // still skip over it and find the next matching line without losing data.
        let mut buf = Vec::new();
        buf.resize(9000, b'@');
        buf.extend_from_slice(b"\nRss: 1234 kB\n");
        scanner.reset(buf.as_slice());

        assert_eq!(
            scanner.next_matching_line(prefix).unwrap(),
            Some(b"Rss: 1234 kB".as_ref())
        );
        assert_eq!(scanner.next_matching_line(prefix).unwrap(), None);
    }

    #[test]
    fn parse_kb_value_as_bytes_converts_kilobytes_to_bytes() {
        // The procfs values are in kibibytes, so they are multiplied by 1024 to get bytes.
        assert_eq!(super::parse_kb_value_as_bytes(b"1234 kB"), Some(1234 * 1024));
        assert_eq!(super::parse_kb_value_as_bytes(b"0 kB"), Some(0));

        // Without a trailing space there is no delimiter, so no value can be extracted.
        assert_eq!(super::parse_kb_value_as_bytes(b"1234"), None);

        // A non-numeric value fails to parse.
        assert_eq!(super::parse_kb_value_as_bytes(b"abc kB"), None);
    }

    #[test]
    fn smaps_rollup_uses_the_first_rss_line() {
        // smaps_rollup is pre-aggregated: the first "Rss:" line is the whole-process RSS, converted from kB to bytes.
        let mut scanner = super::Scanner::new();
        scanner.reset(b"Pss: 10 kB\nRss: 512 kB\nRss: 999 kB\n".as_slice());
        assert_eq!(super::first_rss_value(&mut scanner), Some(512 * 1024));
    }

    #[test]
    fn smaps_sums_every_rss_line() {
        // smaps has one "Rss:" line per mapping, so the process RSS is the sum of them all (in bytes).
        let mut scanner = super::Scanner::new();
        scanner.reset(b"Rss: 4 kB\nPss: 100 kB\nRss: 8 kB\nRss: 16 kB\n".as_slice());
        assert_eq!(super::sum_rss_values(&mut scanner), Some((4 + 8 + 16) * 1024));
    }

    #[test]
    fn smaps_strategies_return_none_when_no_rss_lines_present() {
        // The documented `None`-return path: nothing matched, so neither strategy can report a value.
        let mut scanner = super::Scanner::new();
        scanner.reset(b"Pss: 4 kB\nShared_Clean: 8 kB\n".as_slice());
        assert_eq!(super::first_rss_value(&mut scanner), None);

        let mut scanner = super::Scanner::new();
        scanner.reset(b"Pss: 4 kB\nShared_Clean: 8 kB\n".as_slice());
        assert_eq!(super::sum_rss_values(&mut scanner), None);
    }

    #[test]
    fn parse_statm_rss_bytes_multiplies_resident_pages_by_page_size() {
        // statm fields are "size resident shared text lib data dt"; resident (the second field) is in pages.
        assert_eq!(
            super::parse_statm_rss_bytes(b"1000 42 5 1 0 3 0", 4096),
            Some(42 * 4096)
        );

        // A missing resident field, or a non-numeric one, yields `None`.
        assert_eq!(super::parse_statm_rss_bytes(b"1000", 4096), None);
        assert_eq!(super::parse_statm_rss_bytes(b"1000 xyz", 4096), None);
    }

    #[test]
    fn stat_source_selection_follows_documented_fallback_order() {
        use super::{select_stat_source, StatSource};

        // smaps_rollup is preferred whenever it is available, regardless of the others.
        assert!(matches!(
            select_stat_source(true, true, Some(4096)),
            StatSource::SmapsRollup(_)
        ));
        assert!(matches!(
            select_stat_source(true, false, None),
            StatSource::SmapsRollup(_)
        ));

        // smaps is used only when smaps_rollup is unavailable.
        assert!(matches!(
            select_stat_source(false, true, Some(4096)),
            StatSource::Smaps(_)
        ));

        // statm is the final fallback, carrying whatever page size was resolved (including `None`).
        assert!(matches!(
            select_stat_source(false, false, Some(4096)),
            StatSource::Statm(Some(4096))
        ));
        assert!(matches!(
            select_stat_source(false, false, None),
            StatSource::Statm(None)
        ));
    }
}