cargo_zng/
trace.rs

1use std::{
2    io::{Read as _, Write},
3    path::PathBuf,
4    time::SystemTime,
5};
6
7use clap::*;
8use serde::Deserialize as _;
9
10#[derive(Args, Debug, Default)]
11pub struct TraceArgs {
12    /// Path or command to run the Zng executable
13    ///
14    /// Example: `cargo zng "./some/exe"` or `cargo zng -- cargo run exe`
15    #[arg(trailing_var_arg = true)]
16    command: Vec<String>,
17
18    /// env_logger style filter
19    #[arg(long, short, default_value = "trace")]
20    filter: String,
21
22    /// Output JSON file
23    ///
24    /// {timestamp} and {ts} is replaced with a timestamp in microseconds from Unix epoch
25    #[arg(long, short, default_value = "./trace-{timestamp}.json")]
26    output: String,
27}
28
29pub fn run(args: TraceArgs) {
30    let mut cmd = {
31        let mut cmd = args.command.into_iter().peekable();
32        if let Some(c) = cmd.peek()
33            && c == "--"
34        {
35            cmd.next();
36        }
37        if let Some(c) = cmd.next() {
38            let mut o = std::process::Command::new(c);
39            o.args(cmd);
40            o
41        } else {
42            fatal!("COMMAND is required")
43        }
44    };
45
46    let ts = SystemTime::now()
47        .duration_since(SystemTime::UNIX_EPOCH)
48        .unwrap_or_default()
49        .as_micros()
50        .to_string();
51
52    let tmp = std::env::temp_dir().join("cargo-zng-trace");
53    if let Err(e) = std::fs::create_dir_all(&tmp) {
54        fatal!("cannot create temp dir, {e}");
55    }
56    let out_dir = tmp.join(&ts);
57    let _ = std::fs::remove_dir_all(&out_dir);
58
59    let out_file = PathBuf::from(args.output.replace("{timestamp}", &ts).replace("{ts}", &ts));
60    if let Some(p) = out_file.parent()
61        && let Err(e) = std::fs::create_dir_all(p)
62    {
63        fatal!("cannot output to {}, {e}", out_file.display());
64    }
65    let mut out = match std::fs::File::create(&out_file) {
66        Ok(f) => f,
67        Err(e) => fatal!("cannot output to {}, {e}", out_file.display()),
68    };
69
70    cmd.env("ZNG_RECORD_TRACE", "")
71        .env("ZNG_RECORD_TRACE_DIR", &tmp)
72        .env("ZNG_RECORD_TRACE_FILTER", args.filter)
73        .env("ZNG_RECORD_TRACE_TIMESTAMP", &ts);
74
75    let mut cmd = match cmd.spawn() {
76        Ok(c) => c,
77        Err(e) => fatal!("cannot run, {e}"),
78    };
79
80    let code = match cmd.wait() {
81        Ok(s) => s.code().unwrap_or(0),
82        Err(e) => {
83            error!("cannot wait command exit, {e}");
84            101
85        }
86    };
87
88    if !out_dir.exists() {
89        fatal!("run did not save any trace\nnote: the feature \"trace_recorder\" must be enabled during build")
90    }
91
92    println!("merging trace files...");
93
94    out.write_all(b"[\n")
95        .unwrap_or_else(|e| fatal!("cannot write {}, {e}", out_file.display()));
96    let mut separator = "";
97
98    for trace in glob::glob(out_dir.join("*.json").display().to_string().as_str())
99        .ok()
100        .into_iter()
101        .flatten()
102    {
103        let trace = match trace {
104            Ok(t) => t,
105            Err(e) => {
106                error!("error globing trace files, {e}");
107                continue;
108            }
109        };
110        let json = match std::fs::read_to_string(&trace) {
111            Ok(s) => s,
112            Err(e) => {
113                error!("cannot read {}, {e}", trace.display());
114                continue;
115            }
116        };
117
118        let name_sys_pid = trace
119            .file_name()
120            .unwrap_or_default()
121            .to_string_lossy()
122            .strip_suffix(".json")
123            .unwrap_or_default()
124            .to_owned();
125        let name_sys_pid = match name_sys_pid.parse::<u64>() {
126            Ok(i) => i,
127            Err(_) => {
128                error!("expected only {{pid}}.json files");
129                continue;
130            }
131        };
132
133        // skip the array opening
134        let json = json.trim_start();
135        if !json.starts_with('[') {
136            error!("unknown format in {}", trace.display());
137            continue;
138        }
139        let json = &json[1..];
140
141        let mut reader = std::io::Cursor::new(json.as_bytes());
142        loop {
143            // skip white space and commas to the next object
144            let mut pos = reader.position();
145            let mut buf = [0u8];
146            while reader.read(&mut buf).is_ok() {
147                if !b" \r\n\t,".contains(&buf[0]) {
148                    break;
149                }
150                pos = reader.position();
151            }
152            reader.set_position(pos);
153            let mut de = serde_json::Deserializer::from_reader(&mut reader);
154            match serde_json::Value::deserialize(&mut de) {
155                Ok(mut entry) => {
156                    // patch "pid" to be unique
157                    if let Some(serde_json::Value::Number(pid)) = entry.get_mut("pid") {
158                        if pid.as_u64() != Some(1) {
159                            error!("expected only pid:1 in trace file");
160                            continue;
161                        }
162                        *pid = serde_json::Number::from(name_sys_pid);
163                    }
164
165                    // convert the INFO message process name to actual "process_name" metadata
166                    match &entry {
167                        serde_json::Value::Object(entry) => {
168                            if let Some(serde_json::Value::String(ph)) = entry.get("ph")
169                                && ph == "i"
170                                && let Some(serde_json::Value::Object(args)) = entry.get("args")
171                                && let Some(serde_json::Value::String(msg)) = args.get("message")
172                                && let Some(rest) = msg.strip_prefix("pid: ")
173                                && let Some((sys_pid, p_name)) = rest.split_once(", name: ")
174                                && let Ok(sys_pid) = sys_pid.parse::<u64>()
175                                && name_sys_pid == sys_pid
176                            {
177                                out.write_fmt(format_args!(
178                                    r#"{separator}{{"ph":"M","pid":{sys_pid},"name":"process_name","args":{{"name":"{p_name}"}}}}"#,
179                                ))
180                                .unwrap_or_else(|e| fatal!("cannot write {}, {e}", out_file.display()));
181                            }
182                        }
183                        _ => {
184                            error!("unknown format in {}", trace.display());
185                        }
186                    }
187
188                    out.write_all(separator.as_bytes())
189                        .unwrap_or_else(|e| fatal!("cannot write {}, {e}", out_file.display()));
190                    serde_json::to_writer(&mut out, &entry).unwrap_or_else(|e| fatal!("cannot write {}, {e}", out_file.display()));
191                    separator = ",\n";
192                }
193                Err(_) => break,
194            }
195        }
196    }
197
198    out.write_all(b"\n]")
199        .unwrap_or_else(|e| fatal!("cannot write {}, {e}", out_file.display()));
200    println!("saved to {}", out_file.display());
201
202    if code == 0 {
203        crate::util::exit();
204    } else {
205        // forward the exit code from the exe or cmd
206        std::process::exit(code);
207    }
208}
cargo_zng/trace.rs

cargo_zng/
trace.rs
