Detect the greenlet module and try to handle it

Attempt to detect when the `greenlet` module is used (whether it has
already been imported at the point when tracking starts, or it gets
imported after tracking has started). If we discover that the `greenlet`
module is in use, register a callback to discover greenlet switch events
and react to them.  Unfortunately, this callback needs to be registered
once per thread, and the only place where we can reasonably register it
in a background thread is when our profile function is called, so we may
miss greenlet switches if they occur before any Python calls.

Signed-off-by: Matt Wozniski <mwozniski@bloomberg.net>

news/185.feature.rst

@@ -0,0 +1 @@
+Add experimental support for Greenlet.

requirements-test.txt

@@ -1,2 +1,3 @@
+greenlet
 pytest
 pytest-cov

setup.py

@@ -106,6 +106,7 @@ def build_js_files(self):
 
 test_requires = [
     "Cython",
+    "greenlet",
     "pytest",
     "pytest-cov",
 ]

src/memray/_memray.pyi

@@ -166,6 +166,8 @@ class Tracker:
         exctb: Optional[TracebackType],
     ) -> bool: ...
 
+def greenlet_trace(event: str, args: Any) -> None: ...
+
 class PymallocDomain(enum.IntEnum):
     PYMALLOC_RAW: int
     PYMALLOC_MEM: int

src/memray/_memray.pyx

@@ -40,7 +40,9 @@ from _memray.socket_reader_thread cimport BackgroundSocketReader
 from _memray.source cimport FileSource
 from _memray.source cimport SocketSource
 from _memray.tracking_api cimport Tracker as NativeTracker
+from _memray.tracking_api cimport begin_tracking_greenlets
 from _memray.tracking_api cimport forget_python_stack
+from _memray.tracking_api cimport handle_greenlet_switch
 from _memray.tracking_api cimport install_trace_function
 from cpython cimport PyErr_CheckSignals
 from libc.stdint cimport uint64_t
@@ -361,6 +363,9 @@ cdef class Tracker:
         self._previous_thread_profile_func = threading._profile_hook
         threading.setprofile(start_thread_trace)
 
+        if "greenlet._greenlet" in sys.modules:
+            begin_tracking_greenlets()
+
         NativeTracker.createTracker(
             move(writer),
             self._native_traces,
@@ -383,6 +388,11 @@ def start_thread_trace(frame, event, arg):
     return start_thread_trace
 
 
+def greenlet_trace_function(event, args):
+    if event in {"switch", "throw"}:
+        handle_greenlet_switch(args[0], args[1])
+
+
 cdef millis_to_dt(millis):
     return datetime.fromtimestamp(millis // 1000).replace(
         microsecond=millis % 1000 * 1000)

src/memray/_memray/hooks.cpp

@@ -277,7 +277,12 @@ dlopen(const char* filename, int flag) noexcept
     assert(hooks::dlopen);
 
     void* ret = hooks::dlopen(filename, flag);
-    if (ret) tracking_api::Tracker::invalidate_module_cache();
+    if (ret) {
+        tracking_api::Tracker::invalidate_module_cache();
+        if (filename && 0 != strstr(filename, "/_greenlet.")) {
+            tracking_api::begin_tracking_greenlets();
+        }
+    }
     return ret;
 }
 

src/memray/_memray/tracking_api.cpp

@@ -12,6 +12,7 @@
 #    include <mach/task.h>
 #endif
 
+#include <algorithm>
 #include <mutex>
 #include <type_traits>
 #include <unistd.h>
@@ -94,6 +95,7 @@ class PythonStackTracker
     };
 
   public:
+    static bool s_greenlet_tracking_enabled;
     static bool s_native_tracking_enabled;
 
     static void installProfileHooks();
@@ -107,6 +109,9 @@ class PythonStackTracker
     int pushPythonFrame(PyFrameObject* frame);
     void popPythonFrame(PyFrameObject* frame);
 
+    void installGreenletTraceFunctionIfNeeded();
+    void handleGreenletSwitch(PyObject* from, PyObject* to);
+
   private:
     // Fetch the thread-local stack tracker without checking if its stack needs to be reloaded.
     static PythonStackTracker& getUnsafe();
@@ -124,8 +129,10 @@ class PythonStackTracker
     uint32_t d_num_pending_pops{};
     uint32_t d_tracker_generation{};
     std::vector<LazilyEmittedFrame>* d_stack{};
+    bool d_greenlet_hooks_installed{};
 };
 
+bool PythonStackTracker::s_greenlet_tracking_enabled{false};
 bool PythonStackTracker::s_native_tracking_enabled{false};
 
 std::mutex PythonStackTracker::s_mutex;
@@ -252,6 +259,8 @@ PythonStackTracker::reloadStackIfTrackerChanged()
 int
 PythonStackTracker::pushPythonFrame(PyFrameObject* frame)
 {
+    installGreenletTraceFunctionIfNeeded();
+
     PyCodeObject* code = compat::frameGetCode(frame);
     const char* function = PyUnicode_AsUTF8(code->co_name);
     if (function == nullptr) {
@@ -284,6 +293,8 @@ PythonStackTracker::pushLazilyEmittedFrame(const LazilyEmittedFrame& frame)
 void
 PythonStackTracker::popPythonFrame(PyFrameObject* frame)
 {
+    installGreenletTraceFunctionIfNeeded();
+
     // Note: We check if frame == d_stack->back().frame because Cython could
     // have called our tracing function with profiled Cython calls that we
     // later discarded in favor of the interpreter's stack when a new tracker
@@ -300,6 +311,82 @@ PythonStackTracker::popPythonFrame(PyFrameObject* frame)
     invalidateMostRecentFrameLineNumber();
 }
 
+void
+PythonStackTracker::installGreenletTraceFunctionIfNeeded()
+{
+    if (!s_greenlet_tracking_enabled || d_greenlet_hooks_installed) {
+        return;  // Nothing to do.
+    }
+
+    assert(PyGILState_Check());
+
+    RecursionGuard guard;
+
+    // Borrowed reference
+    PyObject* modules = PySys_GetObject("modules");
+    if (!modules) {
+        return;
+    }
+
+    // Borrowed reference
+    // Look directly at `sys.modules` since we only want to do something if
+    // `greenlet._greenlet` has already been imported.
+    PyObject* _greenlet = PyDict_GetItemString(modules, "greenlet._greenlet");
+    if (!_greenlet) {
+        return;
+    }
+
+    // Borrowed reference
+    PyObject* _memray = PyDict_GetItemString(modules, "memray._memray");
+    if (!_memray) {
+        return;
+    }
+
+    PyObject* ret = PyObject_CallMethod(
+            _greenlet,
+            "settrace",
+            "N",
+            PyObject_GetAttrString(_memray, "greenlet_trace_function"));
+    Py_XDECREF(ret);
+
+    if (!ret) {
+        // This might be hit from PyGILState_Ensure when a new thread state is
+        // created on a C thread, so we can't reasonably raise the exception.
+        PyErr_Print();
+        _exit(1);
+    }
+
+    // Note: guarded by the GIL
+    d_greenlet_hooks_installed = true;
+}
+
+void
+PythonStackTracker::handleGreenletSwitch(PyObject*, PyObject*)
+{
+    RecursionGuard guard;
+
+    // Clear any old TLS stack, emitting pops for frames that had been pushed.
+    if (d_stack) {
+        d_num_pending_pops += std::count_if(d_stack->begin(), d_stack->end(), [](const auto& f) {
+            return f.state != FrameState::NOT_EMITTED;
+        });
+        d_stack->clear();
+        emitPendingPushesAndPops();
+    }
+
+    // Re-create our TLS stack from our Python frames, most recent last.
+    // Note: `frame` may be null; the new greenlet may not have a Python stack.
+    PyFrameObject* frame = PyEval_GetFrame();
+
+    std::vector<PyFrameObject*> stack;
+    while (frame) {
+        stack.push_back(frame);
+        frame = compat::frameGetBack(frame);
+    }
+
+    std::for_each(stack.rbegin(), stack.rend(), [this](auto& frame) { pushPythonFrame(frame); });
+}
+
 std::atomic<bool> Tracker::d_active = false;
 std::unique_ptr<Tracker> Tracker::d_instance_owner;
 std::atomic<Tracker*> Tracker::d_instance = nullptr;
@@ -992,6 +1079,19 @@ forget_python_stack()
     PythonStackTracker::get().clear();
 }
 
+void
+begin_tracking_greenlets()
+{
+    assert(PyGILState_Check());
+    PythonStackTracker::s_greenlet_tracking_enabled = true;
+}
+
+void
+handle_greenlet_switch(PyObject* from, PyObject* to)
+{
+    PythonStackTracker::get().handleGreenletSwitch(from, to);
+}
+
 void
 install_trace_function()
 {
@@ -1026,6 +1126,8 @@ install_trace_function()
     for (auto frame_it = stack.rbegin(); frame_it != stack.rend(); ++frame_it) {
         python_stack_tracker.pushPythonFrame(*frame_it);
     }
+
+    python_stack_tracker.installGreenletTraceFunctionIfNeeded();
 }
 
 }  // namespace memray::tracking_api

src/memray/_memray/tracking_api.h

@@ -102,6 +102,18 @@ install_trace_function();
 void
 forget_python_stack();
 
+/**
+ * Sets a flag to enable integration with the `greenlet` module.
+ */
+void
+begin_tracking_greenlets();
+
+/**
+ * Handle a notification of control switching from one greenlet to another.
+ */
+void
+handle_greenlet_switch(PyObject* from, PyObject* to);
+
 class NativeTrace
 {
   public:

src/memray/_memray/tracking_api.pxd

@@ -7,6 +7,8 @@ from libcpp.string cimport string
 cdef extern from "tracking_api.h" namespace "memray::tracking_api":
     void forget_python_stack() except*
     void install_trace_function() except*
+    void begin_tracking_greenlets() except+
+    void handle_greenlet_switch(object, object) except+
 
     cdef cppclass Tracker:
         @staticmethod