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

capture_dumps/
capture-dumps.rs
