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 <[email protected]>
Co-authored-by: Ben Frederickson <[email protected]>

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

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

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 {

src/speedscope.rs

@@ -288,6 +288,7 @@ mod tests {
             short_filename: None,
             line: 0,
             locals: None,
+            is_entry: true,
         };
 
         let trace = stack_trace::StackTrace {

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,13 +160,16 @@ where
             None
         };
 
+        let is_entry = frame.is_entry();
+
         frames.push(Frame {
             name,
             filename,
             line,
             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,
         }
     }
 }