Fix native stack merging for Python 3.12/3.13 (#751)

* use last_is_shim information for merging python frames

Signed-off-by: Kai Fricke <coding@kaifricke.com>

* fix order for shim detection

Signed-off-by: Kai Fricke <coding@kaifricke.com>

* remove version check

Signed-off-by: Kai Fricke <coding@kaifricke.com>

* fix unittest compile

---------

Signed-off-by: Kai Fricke <coding@kaifricke.com>
Co-authored-by: Ben Frederickson <github@benfrederickson.com>

Author: krfricke

src/main.rs

@@ -285,6 +285,7 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
                     line: 0,
                     locals: None,
                     is_entry: true,
+                    is_shim_entry: true,
                 });
             }
 

src/native_stack_trace.rs

@@ -94,11 +94,14 @@ impl NativeStack {
                         // merge it into the stack. (if we're out of bounds a later
                         // check will pick up - and report overall totals mismatch)
 
-                        // Merge all python frames until we hit one with `is_entry`.
+                        // Merge all python frames until we hit one with `is_entry` (py 3.11)
+                        // or `is_entry_shim` (py 3.12+)
                         while python_frame_index < frames.len() {
                             merged.push(frames[python_frame_index].clone());
 
-                            if frames[python_frame_index].is_entry {
+                            if frames[python_frame_index].is_entry
+                                || frames[python_frame_index].is_shim_entry
+                            {
                                 break;
                             }
 
@@ -150,6 +153,7 @@ impl NativeStack {
                         module: None,
                         locals: None,
                         is_entry: true,
+                        is_shim_entry: true,
                     });
                 });
 
@@ -291,6 +295,7 @@ impl NativeStack {
                     module: Some(frame.module.clone()),
                     locals: None,
                     is_entry: true,
+                    is_shim_entry: true,
                 })
             }
             None => Some(Frame {
@@ -301,6 +306,7 @@ impl NativeStack {
                 short_filename: None,
                 module: Some(frame.module.clone()),
                 is_entry: true,
+                is_shim_entry: true,
             }),
         }
     }

src/speedscope.rs

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

src/stack_trace.rs

@@ -47,8 +47,10 @@ 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.
+    /// If this is an entry frame. Each entry frame corresponds to one native frame (Python 3.11)
     pub is_entry: bool,
+    /// If the last frame was a shim. This is used in Python 3.12+ to detect entry frames.
+    pub is_shim_entry: bool,
 }
 
 #[derive(Debug, Hash, Eq, PartialEq, Ord, PartialOrd, Clone, Serialize)]
@@ -124,6 +126,16 @@ where
     }
 
     let mut frame_ptr = thread.frame(frame_address);
+
+    // We are iterating in reverse, i.e. from last call to first call.
+    // Since Python 3.12, there are shim frames inserted before a block
+    // of Python frames. When we encounter one, update the last frame.
+    let set_last_frame_as_shim_entry = &mut |frames: &mut Vec<Frame>| {
+        if let Some(frame) = frames.last_mut() {
+            frame.is_shim_entry = true;
+        }
+    };
+
     while !frame_ptr.is_null() {
         let frame = process
             .copy_pointer(frame_ptr)
@@ -143,8 +155,10 @@ where
         // would also have to figure out what the address of PyCode_Type is (which will be
         // easier if something like https://github.com/python/cpython/issues/100987#issuecomment-1487227139
         // is merged )
+        // Unset file/function name in py3.13 means this is a shim.
         if filename.is_err() || name.is_err() {
             frame_ptr = frame.back();
+            set_last_frame_as_shim_entry(&mut frames);
             continue;
         }
         let filename = filename?;
@@ -153,6 +167,7 @@ where
         // skip <shim> entries in python 3.12+
         if filename == "<shim>" {
             frame_ptr = frame.back();
+            set_last_frame_as_shim_entry(&mut frames);
             continue;
         }
 
@@ -191,6 +206,7 @@ where
             module: None,
             locals,
             is_entry,
+            is_shim_entry: false,
         });
         if frames.len() > 4096 {
             return Err(format_err!("Max frame recursion depth reached"));
@@ -199,6 +215,9 @@ where
         frame_ptr = frame.back();
     }
 
+    // First frame is always a shim
+    set_last_frame_as_shim_entry(&mut frames);
+
     Ok(StackTrace {
         pid: 0,
         frames,
@@ -322,6 +341,7 @@ impl ProcessInfo {
             line: 0,
             locals: None,
             is_entry: true,
+            is_shim_entry: true,
         }
     }
 }