gh-115946: Add a new interface for stack retrieval for external profilers and debuggers

pablogsal · pablogsal · commit 728c3a27ff28 · 2024-02-26T18:48:27.000Z
diff --git a/Include/cpython/pystate.h b/Include/cpython/pystate.h
@@ -73,6 +73,10 @@ struct _ts {
        pycore_ceval.h. */
     uintptr_t eval_breaker;
 
+    Py_ssize_t n_eval_frames;
+    Py_ssize_t eval_stack_size;
+    PyObject** eval_stack;
+
     struct {
         /* Has been initialized to a safe state.
 
diff --git a/Include/internal/pycore_frame.h b/Include/internal/pycore_frame.h
@@ -99,6 +99,18 @@ static inline void _PyFrame_StackPush(_PyInterpreterFrame *f, PyObject *value) {
     f->stacktop++;
 }
 
+static inline void _PyFrame_PushState(PyThreadState *tstate, _PyInterpreterFrame *f) {
+    if (tstate->n_eval_frames >= tstate->eval_stack_size) {
+       tstate->eval_stack_size *= 2;
+       tstate->eval_stack = PyMem_RawRealloc(tstate->eval_stack, tstate->eval_stack_size * sizeof(PyObject*));
+       if (tstate->eval_stack == NULL) {
+           Py_FatalError("Failed to allocate eval stack");
+       }
+    }
+    tstate->eval_stack[tstate->n_eval_frames] = f->f_executable;
+    tstate->n_eval_frames += 1;
+}
+
 #define FRAME_SPECIALS_SIZE ((int)((sizeof(_PyInterpreterFrame)-1)/sizeof(PyObject *)))
 
 static inline int
@@ -270,6 +282,7 @@ _PyFrame_PushUnchecked(PyThreadState *tstate, PyFunctionObject *func, int null_l
     tstate->datastack_top += code->co_framesize;
     assert(tstate->datastack_top < tstate->datastack_limit);
     _PyFrame_Initialize(new_frame, func, NULL, code, null_locals_from);
+    _PyFrame_PushState(tstate, new_frame);
     return new_frame;
 }
 
@@ -294,6 +307,7 @@ _PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int
     frame->instr_ptr = _PyCode_CODE(code);
     frame->owner = FRAME_OWNED_BY_THREAD;
     frame->return_offset = 0;
+    _PyFrame_PushState(tstate, frame);
     return frame;
 }
 
diff --git a/Python/ceval.c b/Python/ceval.c
@@ -1701,6 +1701,8 @@ clear_gen_frame(PyThreadState *tstate, _PyInterpreterFrame * frame)
 void
 _PyEval_FrameClearAndPop(PyThreadState *tstate, _PyInterpreterFrame * frame)
 {
+    tstate->n_eval_frames -= 1;
+    assert(tstate->n_eval_frames >= 0);
     if (frame->owner == FRAME_OWNED_BY_THREAD) {
         clear_thread_frame(tstate, frame);
     }
@@ -1722,6 +1724,7 @@ _PyEvalFramePushAndInit(PyThreadState *tstate, PyFunctionObject *func,
         goto fail;
     }
     _PyFrame_Initialize(frame, func, locals, code, 0);
+    _PyFrame_PushState(tstate, frame);
     if (initialize_locals(tstate, func, frame->localsplus, args, argcount, kwnames)) {
         assert(frame->owner == FRAME_OWNED_BY_THREAD);
         clear_thread_frame(tstate, frame);
diff --git a/Python/pystate.c b/Python/pystate.c
@@ -1362,6 +1362,9 @@ init_threadstate(_PyThreadStateImpl *_tstate,
     }
 
     tstate->_status.initialized = 1;
+    tstate->n_eval_frames = 0;
+    tstate->eval_stack_size = tstate->py_recursion_limit * 2;
+    tstate->eval_stack = NULL;
 }
 
 static void
@@ -1428,7 +1431,6 @@ new_threadstate(PyInterpreterState *interp, int whence)
 
     init_threadstate(tstate, interp, id, whence);
     add_threadstate(interp, (PyThreadState *)tstate, old_head);
-
     HEAD_UNLOCK(runtime);
     if (!used_newtstate) {
         // Must be called with lock unlocked to avoid re-entrancy deadlock.
@@ -1440,6 +1442,15 @@ new_threadstate(PyInterpreterState *interp, int whence)
     _Py_qsbr_register(tstate, interp, qsbr_idx);
 #endif
 
+    // Initialize eval stack
+    PyThreadState* result = (PyThreadState*)(tstate);
+    if (result && result->eval_stack == NULL) {
+        result->eval_stack = PyMem_RawMalloc(result->py_recursion_limit * sizeof(PyObject*));
+        if (result->eval_stack == NULL) {
+            PyMem_RawFree(result);
+            return NULL;
+        }
+    }
     return (PyThreadState *)tstate;
 }
 
@@ -1650,6 +1661,8 @@ tstate_delete_common(PyThreadState *tstate)
     clear_datastack(tstate);
 
     tstate->_status.finalized = 1;
+    tstate->n_eval_frames = -1;
+    PyMem_RawFree(tstate->eval_stack);
 }
 
 static void
diff --git a/Python/stdlib_module_names.h.new b/Python/stdlib_module_names.h.new

Original file line number	Diff line number	Diff line change
`@@ -1701,6 +1701,8 @@ clear_gen_frame(PyThreadState tstate, _PyInterpreterFrame frame)`
`1701`	`1701`	`void`
`1702`	`1702`	`_PyEval_FrameClearAndPop(PyThreadState tstate, _PyInterpreterFrame frame)`
`1703`	`1703`	`{`
	`1704`	`+ tstate->n_eval_frames -= 1;`
	`1705`	`+ assert(tstate->n_eval_frames >= 0);`
`1704`	`1706`	`if (frame->owner == FRAME_OWNED_BY_THREAD) {`
`1705`	`1707`	`clear_thread_frame(tstate, frame);`
`1706`	`1708`	`}`
`@@ -1722,6 +1724,7 @@ _PyEvalFramePushAndInit(PyThreadState tstate, PyFunctionObject func,`
`1722`	`1724`	`goto fail;`
`1723`	`1725`	`}`
`1724`	`1726`	`_PyFrame_Initialize(frame, func, locals, code, 0);`
	`1727`	`+ _PyFrame_PushState(tstate, frame);`
`1725`	`1728`	`if (initialize_locals(tstate, func, frame->localsplus, args, argcount, kwnames)) {`
`1726`	`1729`	`assert(frame->owner == FRAME_OWNED_BY_THREAD);`
`1727`	`1730`	`clear_thread_frame(tstate, frame);`