From cbcfdf09456372800510db84e66f0d0d79252062 Mon Sep 17 00:00:00 2001 From: andrewjcg Date: Fri, 10 Nov 2023 17:23:32 -0500 Subject: [PATCH] Add chrome trace format for recording samples (#627) This adds "chrometrace" as a new format for the "record" command, which serializes samples using chrome trace events: https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU --- src/chrometrace.rs | 122 +++++++++++++++++++++++++++++++++++++++++++++ src/config.rs | 1 + src/main.rs | 21 ++++++++ 3 files changed, 144 insertions(+) create mode 100644 src/chrometrace.rs diff --git a/src/chrometrace.rs b/src/chrometrace.rs new file mode 100644 index 00000000..24463800 --- /dev/null +++ b/src/chrometrace.rs @@ -0,0 +1,122 @@ +use std::cmp::min; +use std::collections::HashMap; +use std::io::Write; +use std::time::Instant; + +use anyhow::Error; +use serde_derive::Serialize; + +use crate::stack_trace::Frame; +use crate::stack_trace::StackTrace; + +#[derive(Clone, Debug, Serialize)] +struct Args { + pub filename: String, + pub line: Option, +} + +#[derive(Clone, Debug, Serialize)] +struct Event { + pub args: Args, + pub cat: String, + pub name: String, + pub ph: String, + pub pid: u64, + pub tid: u64, + pub ts: u64, +} + +pub struct Chrometrace { + events: Vec, + start_ts: Instant, + prev_traces: HashMap, + show_linenumbers: bool, +} + +impl Chrometrace { + pub fn new(show_linenumbers: bool) -> Chrometrace { + Chrometrace { + events: Vec::new(), + start_ts: Instant::now(), + prev_traces: HashMap::new(), + show_linenumbers, + } + } + + // Return whether these frames are similar enough such that we should merge + // them, instead of creating separate events for them. + fn should_merge_frames(&self, a: &Frame, b: &Frame) -> bool { + a.name == b.name && a.filename == b.filename && (!self.show_linenumbers || a.line == b.line) + } + + fn event(&self, trace: &StackTrace, frame: &Frame, phase: &str, ts: u64) -> Event { + Event { + tid: trace.thread_id, + pid: trace.pid as u64, + name: frame.name.to_string(), + cat: "py-spy".to_owned(), + ph: phase.to_owned(), + ts, + args: Args { + filename: frame.filename.to_string(), + line: if self.show_linenumbers { + Some(frame.line as u32) + } else { + None + }, + }, + } + } + + pub fn increment(&mut self, trace: &StackTrace) -> std::io::Result<()> { + let now = self.start_ts.elapsed().as_micros() as u64; + + // Load the previous frames for this thread. + let prev_frames = self + .prev_traces + .remove(&trace.thread_id) + .map(|t| t.frames) + .unwrap_or_default(); + + // Find the index where we first see new frames. + let new_idx = prev_frames + .iter() + .rev() + .zip(trace.frames.iter().rev()) + .position(|(a, b)| !self.should_merge_frames(a, b)) + .unwrap_or(min(prev_frames.len(), trace.frames.len())); + + // Publish end events for the previous frames that got dropped in the + // most recent trace. + for frame in prev_frames.iter().rev().skip(new_idx).rev() { + self.events.push(self.event(trace, frame, "E", now)); + } + + // Publish start events for frames that got added in the most recent + // trace. + for frame in trace.frames.iter().rev().skip(new_idx) { + self.events.push(self.event(trace, frame, "B", now)); + } + + // Save this stack trace for the next iteration. + self.prev_traces.insert(trace.thread_id, trace.clone()); + + Ok(()) + } + + pub fn write(&self, w: &mut dyn Write) -> Result<(), Error> { + let mut events = Vec::new(); + events.extend(self.events.to_vec()); + + // Add end events for any unfinished slices. + let now = self.start_ts.elapsed().as_micros() as u64; + for trace in self.prev_traces.values() { + for frame in &trace.frames { + events.push(self.event(trace, frame, "E", now)); + } + } + + writeln!(w, "{}", serde_json::to_string(&events)?)?; + Ok(()) + } +} diff --git a/src/config.rs b/src/config.rs index 55c2151b..d8c936f9 100644 --- a/src/config.rs +++ b/src/config.rs @@ -67,6 +67,7 @@ pub enum FileFormat { flamegraph, raw, speedscope, + chrometrace, } impl FileFormat { diff --git a/src/main.rs b/src/main.rs index 4b200c21..d2745c79 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,6 +4,7 @@ extern crate anyhow; extern crate log; mod binary_parser; +mod chrometrace; mod config; mod console_viewer; #[cfg(target_os = "linux")] @@ -108,6 +109,15 @@ impl Recorder for flamegraph::Flamegraph { } } +impl Recorder for chrometrace::Chrometrace { + fn increment(&mut self, trace: &StackTrace) -> Result<(), Error> { + Ok(self.increment(trace)?) + } + fn write(&self, w: &mut dyn Write) -> Result<(), Error> { + self.write(w) + } +} + pub struct RawFlamegraph(flamegraph::Flamegraph); impl Recorder for RawFlamegraph { @@ -129,6 +139,9 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error> Some(FileFormat::raw) => Box::new(RawFlamegraph(flamegraph::Flamegraph::new( config.show_line_numbers, ))), + Some(FileFormat::chrometrace) => { + Box::new(chrometrace::Chrometrace::new(config.show_line_numbers)) + } None => return Err(format_err!("A file format is required to record samples")), }; @@ -139,6 +152,7 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error> Some(FileFormat::flamegraph) => "svg", Some(FileFormat::speedscope) => "json", Some(FileFormat::raw) => "txt", + Some(FileFormat::chrometrace) => "json", None => return Err(format_err!("A file format is required to record samples")), }; let local_time = Local::now().to_rfc3339_opts(SecondsFormat::Secs, true); @@ -342,6 +356,13 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error> ); println!("{}You can use the flamegraph.pl script from https://github.com/brendangregg/flamegraph to generate a SVG", lede); } + FileFormat::chrometrace => { + println!( + "{}Wrote chrome trace to '{}'. Samples: {} Errors: {}", + lede, filename, samples, errors + ); + println!("{}Visit chrome://tracing to view", lede); + } }; Ok(())