From 2f8cfdd6d62ad3031377ec13f8dbf966fad697e4 Mon Sep 17 00:00:00 2001
From: Kai Fricke <krfricke@users.noreply.github.com>
Date: Fri, 17 Nov 2023 21:22:52 +0000
Subject: [PATCH] Fix Python 3.11 native sampling (#635)

In Python 3.11, merging native frames works a bit differently. We have to detect if a frame is an "entry frame". Once we hit a native frame that corresponds to python, we need to keep merging python frames until we hit the entry frame.

Co-authored-by: Ben Frederickson <github@benfrederickson.com>
Co-authored-by: Ben Frederickson <ben@benfrederickson.com>
---
 src/main.rs                |  1 +
 src/native_stack_trace.rs  | 13 ++++++++++++-
 src/python_interpreters.rs |  7 +++++++
 src/speedscope.rs          |  1 +
 src/stack_trace.rs         |  6 ++++++
 5 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/main.rs b/src/main.rs
index d2745c79..f965cc71 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -284,6 +284,7 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
                     short_filename: None,
                     line: 0,
                     locals: None,
+                    is_entry: true,
                 });
             }
 
diff --git a/src/native_stack_trace.rs b/src/native_stack_trace.rs
index 0e2e1af9..ba7e6c64 100644
--- a/src/native_stack_trace.rs
+++ b/src/native_stack_trace.rs
@@ -93,8 +93,16 @@ impl NativeStack {
                         // if we have a corresponding python frame for the evalframe
                         // merge it into the stack. (if we're out of bounds a later
                         // check will pick up - and report overall totals mismatch)
-                        if python_frame_index < frames.len() {
+
+                        // Merge all python frames until we hit one with `is_entry`.
+                        while python_frame_index < frames.len() {
                             merged.push(frames[python_frame_index].clone());
+
+                            if frames[python_frame_index].is_entry {
+                                break;
+                            }
+
+                            python_frame_index += 1;
                         }
                         python_frame_index += 1;
                     }
@@ -141,6 +149,7 @@ impl NativeStack {
                         short_filename: None,
                         module: None,
                         locals: None,
+                        is_entry: true,
                     });
                 });
 
@@ -281,6 +290,7 @@ impl NativeStack {
                     short_filename: None,
                     module: Some(frame.module.clone()),
                     locals: None,
+                    is_entry: true,
                 })
             }
             None => Some(Frame {
@@ -290,6 +300,7 @@ impl NativeStack {
                 line: 0,
                 short_filename: None,
                 module: Some(frame.module.clone()),
+                is_entry: true,
             }),
         }
     }
diff --git a/src/python_interpreters.rs b/src/python_interpreters.rs
index f2485dee..cae9e85d 100644
--- a/src/python_interpreters.rs
+++ b/src/python_interpreters.rs
@@ -47,6 +47,7 @@ pub trait FrameObject {
     fn code(&self) -> *mut Self::CodeObject;
     fn lasti(&self) -> i32;
     fn back(&self) -> *mut Self;
+    fn is_entry(&self) -> bool;
 }
 
 pub trait CodeObject {
@@ -157,6 +158,9 @@ macro_rules! PythonCommonImpl {
             fn back(&self) -> *mut Self {
                 self.f_back
             }
+            fn is_entry(&self) -> bool {
+                true
+            }
         }
 
         impl Object for $py::PyObject {
@@ -357,6 +361,9 @@ impl FrameObject for v3_11_0::_PyInterpreterFrame {
     fn back(&self) -> *mut Self {
         self.previous
     }
+    fn is_entry(&self) -> bool {
+        self.is_entry
+    }
 }
 
 impl Object for v3_11_0::PyObject {
diff --git a/src/speedscope.rs b/src/speedscope.rs
index faf57a0d..3cc0725c 100644
--- a/src/speedscope.rs
+++ b/src/speedscope.rs
@@ -288,6 +288,7 @@ mod tests {
             short_filename: None,
             line: 0,
             locals: None,
+            is_entry: true,
         };
 
         let trace = stack_trace::StackTrace {
diff --git a/src/stack_trace.rs b/src/stack_trace.rs
index 4eb76623..a0d45632 100644
--- a/src/stack_trace.rs
+++ b/src/stack_trace.rs
@@ -47,6 +47,8 @@ pub struct Frame {
     pub line: i32,
     /// Local Variables associated with the frame
     pub locals: Option<Vec<LocalVariable>>,
+    /// If this is an entry frame. Each entry frame corresponds to one native frame.
+    pub is_entry: bool,
 }
 
 #[derive(Debug, Hash, Eq, PartialEq, Ord, PartialOrd, Clone, Serialize)]
@@ -158,6 +160,8 @@ where
             None
         };
 
+        let is_entry = frame.is_entry();
+
         frames.push(Frame {
             name,
             filename,
@@ -165,6 +169,7 @@ where
             short_filename: None,
             module: None,
             locals,
+            is_entry,
         });
         if frames.len() > 4096 {
             return Err(format_err!("Max frame recursion depth reached"));
@@ -278,6 +283,7 @@ impl ProcessInfo {
             short_filename: None,
             line: 0,
             locals: None,
+            is_entry: true,
         }
     }
 }