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

@@ -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<u32>,
+}
+
+#[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<Event>,
+    start_ts: Instant,
+    prev_traces: HashMap<u64, StackTrace>,
+    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(())
+    }
+}

src/config.rs

@@ -67,6 +67,7 @@ pub enum FileFormat {
     flamegraph,
     raw,
     speedscope,
+    chrometrace,
 }
 
 impl FileFormat {

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(())