Skip to content

Commit

Permalink
Merge branch 'master' into page_align
Browse files Browse the repository at this point in the history
  • Loading branch information
benfred authored Nov 10, 2023
2 parents 016cafe + bb0e7e8 commit fa9cf8f
Show file tree
Hide file tree
Showing 8 changed files with 211 additions and 15 deletions.
122 changes: 122 additions & 0 deletions src/chrometrace.rs
Original file line number Diff line number Diff line change
@@ -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(())
}
}
1 change: 1 addition & 0 deletions src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ pub enum FileFormat {
flamegraph,
raw,
speedscope,
chrometrace,
}

impl FileFormat {
Expand Down
1 change: 1 addition & 0 deletions src/coredump.rs
Original file line number Diff line number Diff line change
Expand Up @@ -441,6 +441,7 @@ mod test {
minor: 9,
patch: 13,
release_flags: "".to_owned(),
build_metadata: None,
};
let python_core = PythonCoreDump {
core,
Expand Down
21 changes: 21 additions & 0 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ extern crate anyhow;
extern crate log;

mod binary_parser;
mod chrometrace;
mod config;
mod console_viewer;
#[cfg(target_os = "linux")]
Expand Down Expand Up @@ -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 {
Expand All @@ -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")),
};

Expand All @@ -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);
Expand Down Expand Up @@ -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(())
Expand Down
5 changes: 4 additions & 1 deletion src/python_bindings/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions src/python_process_info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,7 @@ where
minor,
patch: 0,
release_flags: "".to_owned(),
build_metadata: None,
});
}
}
Expand Down
29 changes: 23 additions & 6 deletions src/python_spy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
}
Expand Down
Loading

0 comments on commit fa9cf8f

Please sign in to comment.