Skip to content

Commit

Permalink
Protect against deregistered profile functions in greenlet switches
Browse files Browse the repository at this point in the history
When greenlet tracking is enabled is possible that we run into a
situation where the function that recreates the Python stack in our TLS
variable after a greenlet switch is called **after** the profile
function has been deactivated. In this case, recreating the Python stack
is wrong as we are no longer tracking POP/PUSH events so when the stack
is inspected later nothing guarantees that the frames are still valid.
  • Loading branch information
pablogsal committed Nov 27, 2024
1 parent f5eb8d7 commit 06ddc45
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions src/memray/_memray/tracking_api.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,9 @@ PythonStackTracker::handleGreenletSwitch(PyObject* from, PyObject* to)
// Note: `frame` may be null; the new greenlet may not have a Python stack.
PyFrameObject* frame = PyEval_GetFrame();

// Print refcount of frame
_PyObject_Dump((PyObject*)frame);

std::vector<PyFrameObject*> stack;
while (frame) {
stack.push_back(frame);
Expand Down Expand Up @@ -1219,6 +1222,16 @@ Tracker::handleGreenletSwitch(PyObject* from, PyObject* to)
// Grab the Tracker lock, as this may need to write pushes/pops.
std::unique_lock<std::mutex> lock(*s_mutex);
RecursionGuard guard;

// Check if the trace function is still installed in the current thread as
// it's possible that the thread is dying and the trace function has been
// uninstalled and therefore the stack will not be properly tracked so we
// should not recreate the current thread stack.
PyThreadState* ts = PyThreadState_Get();
if (ts->c_profilefunc != PyTraceFunction) {
return;
}

PythonStackTracker::get().handleGreenletSwitch(from, to);
}

Expand Down

0 comments on commit 06ddc45

Please sign in to comment.