From 06ddc4520720e8868e7824379f932bc8f2666024 Mon Sep 17 00:00:00 2001 From: Pablo Galindo Date: Wed, 27 Nov 2024 00:11:07 +0000 Subject: [PATCH] Protect against deregistered profile functions in greenlet switches 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. --- src/memray/_memray/tracking_api.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/memray/_memray/tracking_api.cpp b/src/memray/_memray/tracking_api.cpp index a6fa414e3c..2c05c37347 100644 --- a/src/memray/_memray/tracking_api.cpp +++ b/src/memray/_memray/tracking_api.cpp @@ -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 stack; while (frame) { stack.push_back(frame); @@ -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 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); }