From 790538ca04a4e876f3b40bbc30b2badabb5e84bd Mon Sep 17 00:00:00 2001
From: Kai Fricke <kai.fricke@openai.com>
Date: Thu, 16 Nov 2023 15:35:08 +0000
Subject: [PATCH] add stack

Signed-off-by: Kai Fricke <kai.fricke@openai.com>
---
 src/config.rs             |  4 ++--
 src/native_stack_trace.rs | 42 +++++++++++++++++++++++++++++++++++----
 src/python_spy.rs         |  2 +-
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 80140d71..a0a3b5d8 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -165,13 +165,13 @@ impl Config {
             .takes_value(true);
 
         #[cfg(unwind)]
-            let native = Arg::new("native")
+        let native = Arg::new("native")
             .short('n')
             .long("native")
             .help("Collect stack traces from native extensions written in Cython, C or C++");
 
         #[cfg(unwind)]
-            let native_all = Arg::new("native-all")
+        let native_all = Arg::new("native-all")
             .short('N')
             .long("native-all")
             .help("Collect stack traces from native-only threads. Implies `--native`.");
diff --git a/src/native_stack_trace.rs b/src/native_stack_trace.rs
index 19b73972..7beda7de 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 {
@@ -242,8 +242,42 @@ impl NativeStack {
         }
     }
 
-    pub fn add_native_only_threads(&self, process: &remoteprocess::Process, traces: &Vec<&remoteprocess::StackFrame>) {
-        todo!()
+    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
diff --git a/src/python_spy.rs b/src/python_spy.rs
index e6d6d7f9..2901895d 100644
--- a/src/python_spy.rs
+++ b/src/python_spy.rs
@@ -337,7 +337,7 @@ impl PythonSpy {
         #[cfg(unwind)]
         if self.config.native_all {
             if let Some(native) = self.native.as_mut() {
-                native.add_native_only_threads(&self.process, &traces);
+                native.add_native_only_threads(&self.process, &mut traces)?;
             }
         }