1 // Copyright 2023 The ChromiumOS Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 mod record;
6
7 use std::collections::BTreeMap;
8 use std::fs;
9 use std::path::PathBuf;
10 use std::process::Command;
11 use std::sync::mpsc::channel;
12 use std::time::Duration;
13
14 use anyhow::anyhow;
15 use anyhow::bail;
16 use anyhow::Context;
17 use anyhow::Result;
18 use argh::FromArgs;
19 use log::info;
20 use record::*;
21 use regex::Regex;
22
23 // Utilities
24
parse_smaps_rollup(smaps: &str) -> Result<BTreeMap<String, u64>>25 fn parse_smaps_rollup(smaps: &str) -> Result<BTreeMap<String, u64>> {
26 let re = Regex::new(&format!(r"\n([^:]+):\s*(\d+) kB")).unwrap();
27 let mut mp = BTreeMap::new();
28 for (_, [key, val]) in re.captures_iter(smaps).map(|c| c.extract()) {
29 mp.insert(key.to_string(), val.parse::<u64>().unwrap());
30 }
31 Ok(mp)
32 }
33
extract_socket_path(cmdline: &str) -> Result<String>34 fn extract_socket_path(cmdline: &str) -> Result<String> {
35 let re = Regex::new(r"(--socket|-s)(\s+|=)(.*?) ").unwrap();
36 let sock = re
37 .captures(cmdline)
38 .with_context(|| anyhow!("regex didn't match: {cmdline}"))?
39 .get(3)
40 .unwrap()
41 .as_str()
42 .to_string();
43 Ok(sock)
44 }
45
46 #[derive(Debug)]
47 struct Process {
48 pid: u32,
49 }
50
51 impl Process {
new(pid: u32) -> Self52 fn new(pid: u32) -> Self {
53 Self { pid }
54 }
55
name(&self) -> Result<String>56 fn name(&self) -> Result<String> {
57 let pid = self.pid;
58 let name = fs::read_to_string(format!("/proc/{pid}/comm"))?;
59 Ok(name.trim().to_string())
60 }
61
mem_stats(&self) -> Result<ProcRecord>62 fn mem_stats(&self) -> Result<ProcRecord> {
63 let pid = self.pid;
64 let name = self.name()?;
65 let content = fs::read_to_string(format!("/proc/{pid}/smaps_rollup"))?;
66
67 let smaps = parse_smaps_rollup(&content)?;
68
69 Ok(ProcRecord { pid, name, smaps })
70 }
71 }
72
73 #[derive(Debug)]
74 struct CrosvmProc {
75 proc: Process,
76 path: PathBuf,
77 sock: String,
78 child_procs: Vec<Process>,
79 }
80
81 impl CrosvmProc {
new() -> Result<Self>82 fn new() -> Result<Self> {
83 let output = Command::new("pgrep").args(["crosvm$"]).output()?.stdout;
84 let s = std::str::from_utf8(&output)?.trim();
85 if s.contains("\n") {
86 // TODO: Support multiple VMs
87 bail!("multiple crosvm instances are running: {s}");
88 }
89 if s.is_empty() {
90 bail!("no crosvm process found");
91 }
92 let pid = s.parse::<u32>().context("failed to parse crosvm pid")?;
93
94 let cmdline = fs::read_to_string(format!("/proc/{pid}/cmdline"))?.replace("\0", " ");
95 let sock = extract_socket_path(&cmdline).context("failed to extract socket path")?;
96 let path = fs::read_link(format!("/proc/{pid}/exe"))?;
97
98 let mut ret = Self {
99 proc: Process::new(pid),
100 child_procs: vec![],
101 sock,
102 path,
103 };
104 ret.update_children()
105 .context("failed to update child PIDs")?;
106 Ok(ret)
107 }
108
update_children(&mut self) -> Result<()>109 fn update_children(&mut self) -> Result<()> {
110 let pid = self.proc.pid;
111 let output = Command::new("pgrep")
112 .args(["-P", &pid.to_string()])
113 .output()?
114 .stdout;
115 let child_procs = std::str::from_utf8(&output)?
116 .trim()
117 .split("\n")
118 .filter(|s| !s.is_empty())
119 .map(|s| s.parse::<u32>().map(Process::new))
120 .collect::<std::result::Result<Vec<_>, _>>()?;
121 self.child_procs = child_procs;
122 Ok(())
123 }
124
balloon_stats(&self) -> Result<BalloonStats>125 fn balloon_stats(&self) -> Result<BalloonStats> {
126 let output = Command::new(self.path.to_str().unwrap())
127 .args(["balloon_stats", &self.sock])
128 .output()?
129 .stdout;
130 let s = std::str::from_utf8(&output)?.trim();
131 let stats: BTreeMap<String, BalloonStats> = serde_json::from_str(s)?;
132 Ok(stats.get("BalloonStats").unwrap().to_owned())
133 }
134
get_record(&mut self, timestamp: u64) -> Result<Record>135 fn get_record(&mut self, timestamp: u64) -> Result<Record> {
136 self.update_children()?;
137
138 let mut stats = vec![];
139 stats.push(self.proc.mem_stats()?);
140 for pid in &self.child_procs {
141 stats.push(pid.mem_stats()?);
142 }
143
144 let balloon_stats = self.balloon_stats().ok();
145
146 Ok(Record {
147 timestamp,
148 stats,
149 balloon_stats,
150 })
151 }
152 }
153
154 #[derive(FromArgs)]
155 /// Argument
156 struct Args {
157 /// duration in second.
158 /// If it's not specified, it runs until Ctrl-C is sent.
159 #[argh(option, short = 'd')]
160 duration: Option<u64>,
161 /// output JSON file path.
162 #[argh(option, short = 'o')]
163 output: String,
164 }
165
wait_for_crosvm() -> CrosvmProc166 fn wait_for_crosvm() -> CrosvmProc {
167 let interval = Duration::from_millis(100);
168 let mut cnt = 0;
169 loop {
170 match CrosvmProc::new() {
171 Ok(crosvm) => {
172 return crosvm;
173 }
174 Err(e) => {
175 if cnt % 10 == 0 {
176 info!("waiting for crosvm starting: {:#}", e);
177 }
178 }
179 }
180 cnt += 1;
181 std::thread::sleep(interval);
182 }
183 }
184
main() -> Result<()>185 fn main() -> Result<()> {
186 env_logger::Builder::from_default_env()
187 .filter(None, log::LevelFilter::Info)
188 .init();
189
190 let args: Args = argh::from_env();
191
192 let mut crosvm = wait_for_crosvm();
193 info!("crosvm process found");
194
195 let mut stats = vec![];
196
197 let start_time = std::time::Instant::now();
198
199 let (tx, rx) = channel();
200 ctrlc::set_handler(move || {
201 println!("Ctrl-C is pressed");
202 tx.send(()).expect("Could not send signal on channel.")
203 })
204 .expect("Error setting Ctrl-C handler");
205
206 let timeout = match args.duration {
207 Some(sec) => {
208 info!("Collect data for {sec} seconds (or until Ctrl-C is sent)");
209 sec
210 }
211 None => {
212 info!("Collect data until Ctrl-C is sent");
213 u64::MAX
214 }
215 };
216
217 for ts in 0..timeout {
218 if let Ok(()) = rx.try_recv() {
219 println!("stop recording");
220 break;
221 }
222 info!("timestamp: {ts} seconds");
223 let rec = match crosvm.get_record(ts) {
224 Ok(r) => r,
225 Err(_) => {
226 info!("crosvm process has gone");
227 break;
228 }
229 };
230 stats.push(rec);
231
232 let now = std::time::Instant::now();
233 let dur = start_time + Duration::from_secs(ts + 1) - now;
234 std::thread::sleep(dur);
235 }
236
237 let json = serde_json::to_string(&stats)?;
238 let result_path = args.output;
239 std::fs::write(&result_path, json)?;
240 println!("Wrote results to {result_path}");
241
242 Ok(())
243 }
244
245 #[cfg(test)]
246 mod tests {
247 use super::*;
248
249 #[test]
test_parse_smaps_rollup()250 fn test_parse_smaps_rollup() {
251 let smaps = r"5561ed990000-ffffffffff601000 ---p 00000000 00:00 0 [rollup]
252 Rss: 391088 kB
253 Pss: 380165 kB
254 Pss_Anon: 270 kB
255 Pss_File: 1350 kB
256 Pss_Shmem: 378543 kB
257 Shared_Clean: 4016 kB
258 Shared_Dirty: 14788 kB
259 Private_Clean: 628 kB
260 Private_Dirty: 371656 kB
261 Referenced: 389220 kB
262 Anonymous: 344 kB
263 LazyFree: 0 kB
264 AnonHugePages: 0 kB
265 ShmemPmdMapped: 0 kB
266 Shared_Hugetlb: 0 kB
267 Private_Hugetlb: 0 kB
268 Swap: 836004 kB
269 SwapPss: 396 kB
270 Locked: 0 kB
271 ";
272
273 let m = parse_smaps_rollup(smaps).unwrap();
274 assert_eq!(m.get("Rss"), Some(&391088));
275 assert_eq!(m.get("Pss"), Some(&380165));
276 assert_eq!(m.get("Pss_Anon"), Some(&270));
277 assert_eq!(m.get("Pss_Shmem"), Some(&378543));
278 assert_eq!(m.get("Shared_Clean"), Some(&4016));
279 assert_eq!(m.get("Shared_Dirty"), Some(&14788));
280 assert_eq!(m.get("Locked"), Some(&0));
281 }
282
283 #[test]
test_extract_socket_path()284 fn test_extract_socket_path() {
285 const EXPECTED_SOCK_PATH: &str = "/path/to/crosvm.sock";
286 let cmd = format!("crosvm run --socket {EXPECTED_SOCK_PATH} vmlinux");
287 assert_eq!(extract_socket_path(&cmd).unwrap(), EXPECTED_SOCK_PATH);
288
289 let cmd = format!("crosvm run --socket={EXPECTED_SOCK_PATH} vmlinux");
290 assert_eq!(extract_socket_path(&cmd).unwrap(), EXPECTED_SOCK_PATH);
291
292 let cmd = format!("crosvm run -s {EXPECTED_SOCK_PATH} vmlinux");
293 assert_eq!(extract_socket_path(&cmd).unwrap(), EXPECTED_SOCK_PATH);
294
295 let cmd = format!("crosvm run -s={EXPECTED_SOCK_PATH} vmlinux");
296 assert_eq!(extract_socket_path(&cmd).unwrap(), EXPECTED_SOCK_PATH);
297 }
298 }
299