Merge branch 'master' into page_align

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/coredump.rs

@@ -441,6 +441,7 @@ mod test {
             minor: 9,
             patch: 13,
             release_flags: "".to_owned(),
+            build_metadata: None,
         };
         let python_core = PythonCoreDump {
             core,

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

src/python_bindings/mod.rs

@@ -222,7 +222,10 @@ pub mod pyruntime {
             },
             Version {
                 major: 3, minor: 8, ..
-            } => Some(1368),
+            } => match version.build_metadata.as_deref() {
+                Some("cinder") => Some(1384),
+                _ => Some(1368),
+            },
             Version {
                 major: 3,
                 minor: 9..=10,

src/python_process_info.rs

@@ -318,6 +318,7 @@ where
                             minor,
                             patch: 0,
                             release_flags: "".to_owned(),
+                            build_metadata: None,
                         });
                     }
                 }

src/python_spy.rs

@@ -189,9 +189,14 @@ impl PythonSpy {
     fn _get_stack_traces<I: InterpreterState>(&mut self) -> Result<Vec<StackTrace>, Error> {
         // Query the OS to get if each thread in the process is running or not
         let mut thread_activity = HashMap::new();
-        for thread in self.process.threads()?.iter() {
-            let threadid: Tid = thread.id()?;
-            thread_activity.insert(threadid, thread.active()?);
+        if self.config.gil_only {
+            // Don't need to collect thread activity if we're only getting the
+            // GIL thread: If we're holding the GIL we're by definition active.
+        } else {
+            for thread in self.process.threads()?.iter() {
+                let threadid: Tid = thread.id()?;
+                thread_activity.insert(threadid, thread.active()?);
+            }
         }
 
         // Lock the process if appropriate. Note we have to lock AFTER getting the thread
@@ -224,6 +229,15 @@ impl PythonSpy {
                 .process
                 .copy_pointer(threads)
                 .context("Failed to copy PyThreadState")?;
+            threads = thread.next();
+
+            let python_thread_id = thread.thread_id();
+            let owns_gil = python_thread_id == gil_thread_id;
+
+            if self.config.gil_only && !owns_gil {
+                continue;
+            }
+
             let mut trace = get_stack_trace(
                 &thread,
                 &self.process,
@@ -232,7 +246,6 @@ impl PythonSpy {
             )?;
 
             // Try getting the native thread id
-            let python_thread_id = thread.thread_id();
 
             // python 3.11+ has the native thread id directly on the PyThreadState object,
             // for older versions of python, try using OS specific code to get the native
@@ -255,7 +268,7 @@ impl PythonSpy {
             }
 
             trace.thread_name = self._get_python_thread_name(python_thread_id);
-            trace.owns_gil = trace.thread_id == gil_thread_id;
+            trace.owns_gil = owns_gil;
             trace.pid = self.process.pid;
 
             // Figure out if the thread is sleeping from the OS if possible
@@ -314,7 +327,11 @@ impl PythonSpy {
                 return Err(format_err!("Max thread recursion depth reached"));
             }
 
-            threads = thread.next();
+            if self.config.gil_only {
+                // There's only one GIL thread and we've captured it, so we can
+                // stop now
+                break;
+            }
         }
         Ok(traces)
     }