capture_dumps/
capture-dumps.rs

1use std::{
2    fs,
3    io::{self, prelude::*},
4    path::Path,
5    process::{ChildStdin, ChildStdout, Command, Stdio},
6};
7
8fn dump_process_to_bytearray(pid: u32, output_dir: &Path, output_name: &Path) -> io::Result<u64> {
9    let output = Command::new("gcore")
10        .args(["-a", &pid.to_string()])
11        .output()?;
12
13    if !output.status.success() {
14        return io::Result::Err(io::Error::new(
15            io::ErrorKind::Other,
16            format!("Failed to dump process: {:?}", output),
17        ));
18    }
19
20    let core_path = format!("core.{}", pid);
21    let output_path = output_dir.join(output_name);
22    let len = fs::copy(&core_path, output_path)?;
23    fs::remove_file(&core_path)?;
24    Ok(len)
25}
26
27fn wait_dump_and_continue(
28    stdin: &mut ChildStdin,
29    stdout: &mut ChildStdout,
30    id: u32,
31    base_dir: &Path,
32    name: &Path,
33) -> Result<(), io::Error> {
34    // Read the input from the process until we get the "Waiting for dump..." message
35    // That way we know the process is ready to be dumped, and we don't need to just sleep a fixed
36    // amount of time
37    loop {
38        let mut buf = [0u8; 1024];
39        let read = stdout.read(&mut buf).unwrap();
40
41        if read == 0 {
42            panic!("Process exited unexpectedly");
43        }
44
45        let buf_str = std::str::from_utf8(&buf[..read]).unwrap();
46        if buf_str.contains("Waiting for dump...") {
47            break;
48        }
49    }
50    let dump_size = dump_process_to_bytearray(id, &base_dir.join("output"), name)?;
51    println!("Got memory dump of file size: {}", dump_size);
52
53    stdin.write_all(b".")?;
54    stdin.flush()?;
55
56    Ok(())
57}
58
59fn main() -> io::Result<()> {
60    let args: Vec<String> = std::env::args().collect();
61    if args.len() < 3 {
62        println!("Usage: ./capture_dumps <binary_path> <base_dir>");
63        std::process::exit(1);
64    }
65
66    let binary_path = &args[1];
67    let base_dir: &Path = args[2].as_ref();
68
69    let mut proc = Command::new(binary_path)
70        .arg(base_dir)
71        .stdout(Stdio::piped())
72        .stdin(Stdio::piped())
73        .spawn()?;
74    let id = proc.id();
75    println!("Started memory testing process with PID: {}", id);
76
77    let stdin = proc.stdin.as_mut().expect("Valid stdin");
78    let stdout = proc.stdout.as_mut().expect("Valid stdin");
79
80    wait_dump_and_continue(stdin, stdout, id, base_dir, "initial_dump.bin".as_ref())?;
81    wait_dump_and_continue(stdin, stdout, id, base_dir, "final_dump.bin".as_ref())?;
82
83    // Wait for the process to finish and print the output
84    let output = proc.wait()?;
85    println!("Return code: {}", output);
86
87    std::process::exit(output.code().unwrap_or(1));
88}
capture_dumps/capture-dumps.rs

capture_dumps/
capture-dumps.rs
