ic_host_process/sampling/
mod.rs1use 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#[derive(Debug)]
20pub enum SamplingError {
21 Io(io::Error),
24 InvalidData(&'static str),
26 Command(Box<ToolError>),
29 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
62pub struct Process {
63 pub parent: u32,
65 pub resident_bytes: u64,
67}
68
69pub type Snapshot = BTreeMap<u32, Process>;
71
72#[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 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 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
204pub 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
254pub 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;