diff --git a/src/config.rs b/src/config.rs index d8c936f9..a0a3b5d8 100644 --- a/src/config.rs +++ b/src/config.rs @@ -18,6 +18,9 @@ pub struct Config { /// the native stack traces pub native: bool, + /// Whether or not to also print native-only threads. Using this implies `native`. + pub native_all: bool, + // The following config options only apply when using py-spy as an application #[doc(hidden)] pub command: String, @@ -127,6 +130,7 @@ impl Default for Config { sampling_rate: 100, duration: RecordDuration::Unlimited, native: false, + native_all: false, gil_only: false, include_idle: false, include_thread_ids: false, @@ -166,6 +170,12 @@ impl Config { .long("native") .help("Collect stack traces from native extensions written in Cython, C or C++"); + #[cfg(unwind)] + let native_all = Arg::new("native-all") + .short('N') + .long("native-all") + .help("Collect stack traces from native-only threads. Implies `--native`."); + #[cfg(not(target_os="freebsd"))] let nonblocking = Arg::new("nonblocking") .long("nonblocking") @@ -335,6 +345,13 @@ impl Config { #[cfg(unwind)] let dump = dump.arg(native.clone()); + #[cfg(unwind)] + let record = record.arg(native_all.clone()); + #[cfg(unwind)] + let top = top.arg(native_all.clone()); + #[cfg(unwind)] + let dump = dump.arg(native_all.clone()); + // Nonblocking isn't an option for freebsd, remove #[cfg(not(target_os = "freebsd"))] let record = record.arg(nonblocking.clone()); @@ -430,7 +447,8 @@ impl Config { .map(|p| p.parse().expect("invalid pid")); config.full_filenames = matches.occurrences_of("full_filenames") > 0; if cfg!(unwind) { - config.native = matches.occurrences_of("native") > 0; + config.native_all = matches.occurrences_of("native-all") > 0; + config.native = config.native_all || matches.occurrences_of("native") > 0; } config.capture_output = config.command != "record" || matches.occurrences_of("capture") > 0; diff --git a/src/native_stack_trace.rs b/src/native_stack_trace.rs index ba7e6c64..c3573520 100644 --- a/src/native_stack_trace.rs +++ b/src/native_stack_trace.rs @@ -5,11 +5,11 @@ use std::num::NonZeroUsize; use cpp_demangle::{BorrowedSymbol, DemangleOptions}; use lazy_static::lazy_static; use lru::LruCache; -use remoteprocess::{self, Pid}; +use remoteprocess::{self, Pid, Tid}; use crate::binary_parser::BinaryInfo; use crate::cython; -use crate::stack_trace::Frame; +use crate::stack_trace::{Frame, StackTrace}; use crate::utils::resolve_filename; pub struct NativeStack { @@ -251,6 +251,44 @@ impl NativeStack { } } + pub fn add_native_only_threads( + &mut self, + process: &remoteprocess::Process, + traces: &mut Vec, + ) -> Result<(), Error> { + // Set of all threads we already processed + let seen_threads = + HashSet::::from_iter(traces.iter().map(|t| t.os_thread_id.unwrap_or(0) as Tid)); + + for native_thread in process.threads()?.into_iter() { + let tid = native_thread.id()?; + + if seen_threads.contains(&tid) { + // We've already seen this thread, don't add it again + continue; + } + + // We are reusing the `merge_native_stack` method and just pass an + // empty python stack. + let native_stack = self.get_thread(&native_thread)?; + let python_stack = Vec::new(); + let symbolized_stack = self.merge_native_stack(&python_stack, native_stack)?; + + // Push new stack trace + traces.push(StackTrace { + pid: process.pid, + thread_id: tid.try_into().unwrap_or(0), + thread_name: None, + os_thread_id: tid.try_into().ok(), + active: native_thread.active().unwrap_or(false), + owns_gil: false, + frames: symbolized_stack, + process_info: None, + }); + } + Ok(()) + } + /// translates a native frame into a optional frame. none indicates we should ignore this frame fn translate_native_frame(&self, frame: &remoteprocess::StackFrame) -> Option { match &frame.function { diff --git a/src/python_spy.rs b/src/python_spy.rs index f5e36d68..2901895d 100644 --- a/src/python_spy.rs +++ b/src/python_spy.rs @@ -333,6 +333,14 @@ impl PythonSpy { break; } } + + #[cfg(unwind)] + if self.config.native_all { + if let Some(native) = self.native.as_mut() { + native.add_native_only_threads(&self.process, &mut traces)?; + } + } + Ok(traces) }