Skip to main content

ic_host_process/sampling/
mod.rs

1//! Native process-leader RSS observations, independent of workload policy.
2//! Linux retains page-based /proc measurements; Darwin reports native ps RSS.
3
4use std::collections::{BTreeMap, BTreeSet};
5
6use crate::tool::{OutputLimits, ToolError};
7
8#[cfg(any(target_os = "macos", test))]
9use std::process::Command;
10#[cfg(target_os = "linux")]
11use std::{
12    fs::{self, File},
13    path::Path,
14};
15
16use std::{fmt, io};
17
18/// Failure to observe a native process table or sum its resident bytes.
19#[derive(Debug)]
20pub enum SamplingError {
21    /// Native I/O, stat-record storage/byte admission or stat UTF-8 failed.
22    /// Non-I/O stream refusals retain an artifact error as the I/O error's cause.
23    Io(io::Error),
24    /// A native record was malformed or its byte arithmetic overflowed.
25    InvalidData(&'static str),
26    /// Native table capture, exit or completeness admission failed, with bounded
27    /// output and independent cleanup evidence retained by the capture engine.
28    Command(Box<ToolError>),
29    /// This host has no supported process-table observation backend.
30    UnsupportedHost,
31}
32
33impl fmt::Display for SamplingError {
34    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
35        match self {
36            Self::Io(error) => error.fmt(formatter),
37            Self::InvalidData(message) => formatter.write_str(message),
38            Self::Command(source) => write!(formatter, "process table capture failed: {source}"),
39            Self::UnsupportedHost => formatter.write_str("native sampling requires Linux or macOS"),
40        }
41    }
42}
43
44impl std::error::Error for SamplingError {
45    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
46        match self {
47            Self::Io(error) => Some(error),
48            Self::Command(error) => Some(error.as_ref()),
49            _ => None,
50        }
51    }
52}
53
54impl From<io::Error> for SamplingError {
55    fn from(error: io::Error) -> Self {
56        Self::Io(error)
57    }
58}
59
60/// One process leader observed in a native table.
61#[derive(Clone, Copy, Debug, Eq, PartialEq)]
62pub struct Process {
63    /// Observed parent process identifier.
64    pub parent: u32,
65    /// Native resident memory in bytes; shared pages may occur in multiple processes.
66    pub resident_bytes: u64,
67}
68
69/// Native observations indexed by process-leader identifier.
70pub type Snapshot = BTreeMap<u32, Process>;
71
72/// Measurement-source identity; consumers retain it with their comparisons.
73#[must_use]
74pub const fn source() -> &'static str {
75    if cfg!(target_os = "linux") {
76        "linux-proc-stat"
77    } else if cfg!(target_os = "macos") {
78        "macos-ps-rss"
79    } else {
80        "unsupported"
81    }
82}
83
84#[cfg(any(target_os = "linux", test))]
85fn parse_stat(stat: &str, page_bytes: u64) -> Result<Process, SamplingError> {
86    // comm can contain spaces and closing parentheses. Remaining fields start
87    // at state (3), then PPID (4), and RSS in pages (24).
88    let (_, fields) = stat.rsplit_once(')').ok_or(SamplingError::InvalidData(
89        "missing process command terminator",
90    ))?;
91    let mut fields = fields.split_whitespace();
92    let parent = fields
93        .nth(1)
94        .and_then(|field| field.parse().ok())
95        .ok_or(SamplingError::InvalidData("invalid process parent"))?;
96    let pages: u64 = fields
97        .nth(19)
98        .and_then(|field| field.parse().ok())
99        .ok_or(SamplingError::InvalidData("invalid process RSS pages"))?;
100    let resident_bytes = pages
101        .checked_mul(page_bytes)
102        .ok_or(SamplingError::InvalidData("RSS byte count overflow"))?;
103    Ok(Process {
104        parent,
105        resident_bytes,
106    })
107}
108
109#[cfg(target_os = "linux")]
110fn process_snapshot(
111    proc: &Path,
112    page_bytes: u64,
113    record_bytes: u64,
114) -> Result<Snapshot, SamplingError> {
115    let record_bytes = usize::try_from(record_bytes)
116        .map_err(|_| SamplingError::InvalidData("stat allowance exceeds the host address space"))?;
117    let mut processes = Snapshot::new();
118    for entry in fs::read_dir(proc)? {
119        let entry = entry?;
120        let Some(pid) = entry
121            .file_name()
122            .to_str()
123            .and_then(|name| name.parse::<u32>().ok())
124        else {
125            continue;
126        };
127        let stat = File::open(entry.path().join("stat")).and_then(|file| {
128            ic_host_artifacts::artifact::read_reader(file, record_bytes).map_err(io::Error::from)
129        });
130        match stat {
131            Ok(stat) => {
132                let stat = std::str::from_utf8(&stat)
133                    .map_err(|source| io::Error::new(io::ErrorKind::InvalidData, source))?;
134                processes.insert(pid, parse_stat(stat, page_bytes)?);
135            }
136            Err(error)
137                if matches!(
138                    error.kind(),
139                    io::ErrorKind::NotFound | io::ErrorKind::PermissionDenied
140                ) => {}
141            Err(error) if error.raw_os_error() == Some(rustix::io::Errno::SRCH.raw_os_error()) => {}
142            Err(error) => return Err(error.into()),
143        }
144    }
145    Ok(processes)
146}
147
148#[cfg(any(target_os = "macos", test))]
149fn parse_ps(output: &str) -> Result<Snapshot, SamplingError> {
150    let mut processes = Snapshot::new();
151    for line in output.lines().filter(|line| !line.trim().is_empty()) {
152        let mut fields = line.split_whitespace();
153        let pid = fields
154            .next()
155            .and_then(|field| field.parse::<u32>().ok())
156            .ok_or(SamplingError::InvalidData("invalid ps PID"))?;
157        let parent = fields
158            .next()
159            .and_then(|field| field.parse().ok())
160            .ok_or(SamplingError::InvalidData("invalid ps parent"))?;
161        let rss: u64 = fields
162            .next()
163            .and_then(|field| field.parse().ok())
164            .ok_or(SamplingError::InvalidData("invalid ps RSS"))?;
165        if fields.next().is_some() {
166            return Err(SamplingError::InvalidData("unexpected ps fields"));
167        }
168        // Darwin ps documents RSS in 1024-byte units, not host memory pages.
169        let resident_bytes = rss
170            .checked_mul(1024)
171            .ok_or(SamplingError::InvalidData("ps RSS byte count overflow"))?;
172        if processes
173            .insert(
174                pid,
175                Process {
176                    parent,
177                    resident_bytes,
178                },
179            )
180            .is_some()
181        {
182            return Err(SamplingError::InvalidData(
183                "duplicate process leader in ps output",
184            ));
185        }
186    }
187    Ok(processes)
188}
189
190#[cfg(any(target_os = "macos", test))]
191fn command_snapshot(
192    command: &mut Command,
193    limits: OutputLimits,
194) -> Result<Snapshot, SamplingError> {
195    let output = crate::tool::capture_command(command, limits)
196        .and_then(crate::tool::CompletedExecution::require_success)
197        .and_then(crate::tool::CompletedExecution::require_complete)
198        .map_err(|source| SamplingError::Command(Box::new(source)))?;
199    let output = std::str::from_utf8(&output.stdout)
200        .map_err(|_| SamplingError::InvalidData("ps output is not UTF-8"))?;
201    parse_ps(output)
202}
203
204/// Observe process leaders using Linux /proc or the macOS native ps table.
205///
206/// `record_bytes` bounds each complete Linux stat record; oversized records refuse
207/// before parsing, even if a retained prefix contains plausible numeric fields.
208/// Consumers own sampling cadence,
209/// workload selection and total snapshot memory. `command_limits` selects macOS
210/// stdout/stderr retention and an optional capture deadline; `timeout: None`
211/// imposes no elapsed-time deadline. Linux does not use `command_limits`, and
212/// macOS does not use `record_bytes`.
213/// Processes that disappear or deny stat access during Linux enumeration are skipped.
214/// macOS ps RSS uses 1024-byte units, independently of native memory page size.
215/// The native command's stderr is captured rather than printed. Failed captures
216/// retain bounded output and cleanup evidence in [`SamplingError::Command`].
217/// Successful but truncated output refuses before parsing, even if the retained
218/// prefix contains complete rows. This is not a bound on total parsed table memory.
219///
220/// # Errors
221/// Returns native I/O, capture/exit/truncation, malformed-record or checked-arithmetic
222/// failures. macOS uses the direct-child capture engine; deadlines and cleanup
223/// have that engine's documented synchronous syscall and reaping limitations.
224pub fn snapshot(
225    record_bytes: u64,
226    command_limits: OutputLimits,
227) -> Result<Snapshot, SamplingError> {
228    #[cfg(not(target_os = "linux"))]
229    let _ = record_bytes;
230    #[cfg(not(target_os = "macos"))]
231    let _ = command_limits;
232    #[cfg(target_os = "linux")]
233    {
234        let page_bytes = rustix::param::page_size();
235        let page_bytes = u64::try_from(page_bytes)
236            .ok()
237            .filter(|value| *value > 0)
238            .ok_or(SamplingError::InvalidData("native page size unavailable"))?;
239        process_snapshot(Path::new("/proc"), page_bytes, record_bytes)
240    }
241    #[cfg(target_os = "macos")]
242    {
243        command_snapshot(
244            Command::new("/bin/ps")
245                .args(["-A", "-o", "pid=,ppid=,rss="])
246                .env("LC_ALL", "C"),
247            command_limits,
248        )
249    }
250    #[cfg(not(any(target_os = "linux", target_os = "macos")))]
251    Err(SamplingError::UnsupportedHost)
252}
253
254/// Sum observed resident bytes for a root and its descendants, counting each leader once.
255///
256/// Missing roots yield zero; cycles cannot repeat an observation. This is a sampled
257/// sum of native RSS, not unique physical memory or proof that a process is still live.
258///
259/// # Errors
260/// Returns an arithmetic overflow rather than wrapping the resident byte total.
261pub fn tree_rss(root: u32, processes: &Snapshot) -> Result<(u64, usize), SamplingError> {
262    if !processes.contains_key(&root) {
263        return Ok((0, 0));
264    }
265    let mut children: BTreeMap<u32, Vec<u32>> = BTreeMap::new();
266    for (pid, process) in processes {
267        children.entry(process.parent).or_default().push(*pid);
268    }
269    let mut pending = vec![root];
270    let mut seen = BTreeSet::new();
271    let mut resident = 0_u64;
272    while let Some(pid) = pending.pop() {
273        if !seen.insert(pid) {
274            continue;
275        }
276        resident = resident
277            .checked_add(processes[&pid].resident_bytes)
278            .ok_or(SamplingError::InvalidData("process-tree RSS overflow"))?;
279        if let Some(children) = children.get(&pid) {
280            pending.extend(children);
281        }
282    }
283    Ok((resident, seen.len()))
284}
285
286#[cfg(test)]
287mod tests;