Add --native-all flag that also print the stack of non-python threads

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;

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<StackTrace>,
+    ) -> Result<(), Error> {
+        // Set of all threads we already processed
+        let seen_threads =
+            HashSet::<Tid>::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<Frame> {
         match &frame.function {

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