11#include "Python.h"
2+ #include "ceval.h"
3+ #if PY_MINOR_VERSION < 11
24#include "code.h"
5+ #else
6+ #include "internal/pycore_code.h"
7+ #include "internal/pycore_frame.h"
8+ #endif
9+ #include "frameobject.h"
310#include "object.h"
11+
412#ifndef _GNU_SOURCE
513#define _GNU_SOURCE
614#endif
7- #include "frameobject.h"
815#include <dlfcn.h>
916#include <pthread.h>
1017#include <stdatomic.h>
1623#include <stdbool.h>
1724#include <errno.h>
1825
26+ #if PY_MINOR_VERSION < 9
27+ PyFrameObject * PyFrame_GetBack (PyFrameObject * frame ) {
28+ if (frame -> f_back != NULL ) {
29+ Py_INCREF (frame -> f_back );
30+ }
31+ return frame -> f_back ;
32+ }
33+
34+ PyCodeObject * PyFrame_GetCode (PyFrameObject * frame ) {
35+ Py_INCREF (frame -> f_code );
36+ return frame -> f_code ;
37+ }
38+ #endif
39+
1940// Macro to create the publicly exposed symbol:
2041#ifdef __APPLE__
2142#define SYMBOL_PREFIX (func ) reimplemented_##func
3354#define likely (x ) __builtin_expect(!!(x), 1)
3455#define unlikely (x ) __builtin_expect(!!(x), 0)
3556
36- // Underlying APIs we're wrapping:
37- static void * (* underlying_real_mmap )(void * addr , size_t length , int prot ,
38- int flags , int fd , off_t offset ) = 0 ;
57+ // Underlying APIs we're wrapping:
58+ static void * (* underlying_real_mmap )(void * addr , size_t length , int prot ,
59+ int flags , int fd , off_t offset ) = 0 ;
3960static int (* underlying_real_pthread_create )(pthread_t * thread ,
4061 const pthread_attr_t * attr ,
4162 void * (* start_routine )(void * ),
@@ -118,8 +139,22 @@ static inline int should_track_memory() {
118139 return (likely (initialized ) && atomic_load_explicit (& tracking_allocations , memory_order_acquire ) && !am_i_reentrant ());
119140}
120141
121- // Current thread's Python state:
122- static _Thread_local PyFrameObject * current_frame = NULL ;
142+ // Current thread's Python state; typically only set in C functions where GIL
143+ // might be released.
144+ static _Thread_local int current_line_number = -1 ;
145+
146+ static inline int get_current_line_number () {
147+ if (PyGILState_Check ()) {
148+ PyFrameObject * frame = PyEval_GetFrame ();
149+ if (frame != NULL ) {
150+ return PyFrame_GetLineNumber (frame );
151+ }
152+ }
153+ if (current_line_number != -1 ) {
154+ return current_line_number ;
155+ }
156+ return 0 ;
157+ }
123158
124159// The file and function name responsible for an allocation.
125160struct FunctionLocation {
@@ -200,13 +235,16 @@ static void __attribute__((constructor)) constructor() {
200235 initialized = 1 ;
201236}
202237
203- static void start_call (uint64_t function_id , uint16_t line_number ) {
238+ static void start_call (uint64_t function_id , uint16_t line_number , PyFrameObject * current_frame ) {
204239 if (should_track_memory ()) {
205240 increment_reentrancy ();
206241 uint16_t parent_line_number = 0 ;
207- if (current_frame != NULL && current_frame -> f_back != NULL ) {
208- PyFrameObject * f = current_frame -> f_back ;
209- parent_line_number = PyFrame_GetLineNumber (f );
242+ if (current_frame != NULL ) {
243+ PyFrameObject * parent = PyFrame_GetBack (current_frame );
244+ if (parent != NULL ){
245+ parent_line_number = PyFrame_GetLineNumber (parent );
246+ Py_DECREF (parent );
247+ }
210248 }
211249 pymemprofile_start_call (parent_line_number , function_id , line_number );
212250 decrement_reentrancy ();
@@ -226,39 +264,53 @@ __attribute__((visibility("hidden"))) int
226264fil_tracer (PyObject * obj , PyFrameObject * frame , int what , PyObject * arg ) {
227265 switch (what ) {
228266 case PyTrace_CALL :
229- // Store the current frame, so malloc() can look up line number:
230- current_frame = frame ;
231-
232267 /*
233268 We want an efficient identifier for filename+fuction name. So we register
234269 the function + filename with some Rust code that gives back its ID, and
235270 then store the ID. Due to bad API design, value 0 indicates "no result",
236271 so we actually store the result + 1.
237272 */
273+ current_line_number = frame -> f_lineno ;
238274 uint64_t function_id = 0 ;
239275 assert (extra_code_index != -1 );
240- _PyCode_GetExtra (( PyObject * ) frame -> f_code , extra_code_index ,
241- (void * * )& function_id );
276+ PyCodeObject * code = PyFrame_GetCode ( frame );
277+ _PyCode_GetExtra (( PyObject * ) code , extra_code_index , (void * * )& function_id );
242278 if (function_id == 0 ) {
243279 Py_ssize_t filename_length , function_length ;
244- const char * filename = PyUnicode_AsUTF8AndSize (frame -> f_code -> co_filename ,
280+ const char * filename = PyUnicode_AsUTF8AndSize (code -> co_filename ,
245281 & filename_length );
246- const char * function_name = PyUnicode_AsUTF8AndSize (frame -> f_code -> co_name ,
282+ const char * function_name = PyUnicode_AsUTF8AndSize (code -> co_name ,
247283 & function_length );
248284 increment_reentrancy ();
249285 function_id = pymemprofile_add_function_location (filename , (uint64_t )filename_length , function_name , (uint64_t )function_length );
250286 decrement_reentrancy ();
251- _PyCode_SetExtra ((PyObject * )frame -> f_code , extra_code_index ,
287+ _PyCode_SetExtra ((PyObject * )code , extra_code_index ,
252288 (void * )function_id + 1 );
289+ Py_DECREF (code );
253290 } else {
254291 function_id -= 1 ;
255292 }
256- start_call (function_id , frame -> f_lineno );
293+ start_call (function_id , current_line_number , frame );
257294 break ;
258295 case PyTrace_RETURN :
259296 finish_call ();
260- // We're done with this frame, so set the parent frame:
261- current_frame = frame -> f_back ;
297+ if (frame != NULL ) {
298+ PyFrameObject * parent = PyFrame_GetBack (frame );
299+ if (parent == NULL ) {
300+ current_line_number = -1 ;
301+ } else {
302+ current_line_number = PyFrame_GetLineNumber (parent );
303+ Py_DECREF (parent );
304+ }
305+ }
306+ break ;
307+ case PyTrace_C_CALL :
308+ // C calls might release GIL, in which case they won't change the line
309+ // number, so record it.
310+ current_line_number = PyFrame_GetLineNumber (frame );
311+ break ;
312+ case PyTrace_C_RETURN :
313+ current_line_number = -1 ;
262314 break ;
263315 default :
264316 break ;
@@ -315,20 +367,12 @@ fil_dump_peak_to_flamegraph(const char *path) {
315367
316368// *** End APIs called by Python ***
317369static void add_allocation (size_t address , size_t size ) {
318- uint16_t line_number = 0 ;
319- PyFrameObject * f = current_frame ;
320- if (f != NULL ) {
321- line_number = PyFrame_GetLineNumber (f );
322- }
370+ uint16_t line_number = get_current_line_number ();
323371 pymemprofile_add_allocation (address , size , line_number );
324372}
325373
326374static void add_anon_mmap (size_t address , size_t size ) {
327- uint16_t line_number = 0 ;
328- PyFrameObject * f = current_frame ;
329- if (f != NULL ) {
330- line_number = PyFrame_GetLineNumber (f );
331- }
375+ uint16_t line_number = get_current_line_number ();
332376 pymemprofile_add_anon_mmap (address , size , line_number );
333377}
334378
0 commit comments