feat(profiling): gevent support

We add support for gevent in stack-v2.

Author: P403n1x87

ddtrace/internal/datadog/profiling/stack_v2/CMakeLists.txt

@@ -41,7 +41,7 @@ endif()
 
 # Add echion
 set(ECHION_COMMIT
-    "6ebe7dddb604aa97e89f072c0fc65c9785e023a0" # https://github.com/P403n1x87/echion/commit/6ebe7dddb604aa97e89f072c0fc65c9785e023a0
+    "dcd2a27a219cb2e91e799779e04ba7276ba13a27" # https://github.com/P403n1x87/echion/commit/dcd2a27a219cb2e91e799779e04ba7276ba13a27
     CACHE STRING "Commit hash of echion to use")
 FetchContent_Declare(
     echion

ddtrace/internal/datadog/profiling/stack_v2/include/sampler.hpp

@@ -51,6 +51,9 @@ class Sampler
                       PyObject* _asyncio_eager_tasks);
     void link_tasks(PyObject* parent, PyObject* child);
     void sampling_thread(const uint64_t seq_num);
+    void track_greenlet(uintptr_t greenlet_id, StringTable::Key name, PyObject* frame_cell);
+    void untrack_greenlet(uintptr_t greenlet_id);
+    void link_greenlets(uintptr_t parent, uintptr_t child);
 
     // The Python side dynamically adjusts the sampling rate based on overhead, so we need to be able to update our
     // own intervals accordingly.  Rather than a preemptive measure, we assume the rate is ~fairly stable and just

ddtrace/internal/datadog/profiling/stack_v2/src/sampler.cpp

@@ -2,6 +2,7 @@
 
 #include "thread_span_links.hpp"
 
+#include "echion/greenlets.h"
 #include "echion/interp.h"
 #include "echion/tasks.h"
 #include "echion/threads.h"
@@ -338,3 +339,38 @@ Sampler::link_tasks(PyObject* parent, PyObject* child)
     std::lock_guard<std::mutex> guard(task_link_map_lock);
     task_link_map[child] = parent;
 }
+
+void
+Sampler::track_greenlet(uintptr_t greenlet_id, StringTable::Key name, PyObject* frame_cell)
+{
+    const std::lock_guard<std::mutex> guard(greenlet_info_map_lock);
+
+    auto entry = greenlet_info_map.find(greenlet_id);
+    if (entry != greenlet_info_map.end())
+        // Greenlet is already tracked so we update its info
+        entry->second = std::make_unique<GreenletInfo>(greenlet_id, frame_cell, name);
+    else
+        greenlet_info_map.emplace(greenlet_id, std::make_unique<GreenletInfo>(greenlet_id, frame_cell, name));
+
+    // Update the thread map
+    auto native_id = PyThread_get_thread_native_id();
+    greenlet_thread_map[native_id] = greenlet_id;
+}
+
+void
+Sampler::untrack_greenlet(uintptr_t greenlet_id)
+{
+    const std::lock_guard<std::mutex> guard(greenlet_info_map_lock);
+
+    greenlet_info_map.erase(greenlet_id);
+    greenlet_parent_map.erase(greenlet_id);
+    greenlet_thread_map.erase(greenlet_id);
+}
+
+void
+Sampler::link_greenlets(uintptr_t parent, uintptr_t child)
+{
+    std::lock_guard<std::mutex> guard(greenlet_info_map_lock);
+
+    greenlet_parent_map[child] = parent;
+}

ddtrace/internal/datadog/profiling/stack_v2/src/stack_v2.cpp

@@ -191,6 +191,56 @@ stack_v2_set_adaptive_sampling(PyObject* Py_UNUSED(self), PyObject* args)
     Py_RETURN_NONE;
 }
 
+static PyObject*
+track_greenlet(PyObject* Py_UNUSED(m), PyObject* args)
+{
+    uintptr_t greenlet_id; // map key
+    PyObject* name;
+    PyObject* frame_cell;
+
+    if (!PyArg_ParseTuple(args, "lOO", &greenlet_id, &name, &frame_cell))
+        return NULL;
+
+    StringTable::Key greenlet_name;
+
+    try {
+        greenlet_name = string_table.key(name);
+    } catch (StringTable::Error&) {
+        // We failed to get this task but we keep going
+        PyErr_SetString(PyExc_RuntimeError, "Failed to get greenlet name from the string table");
+        return NULL;
+    }
+
+    Sampler::get().track_greenlet(greenlet_id, greenlet_name, frame_cell);
+
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+untrack_greenlet(PyObject* Py_UNUSED(m), PyObject* args)
+{
+    uintptr_t greenlet_id;
+    if (!PyArg_ParseTuple(args, "l", &greenlet_id))
+        return NULL;
+
+    Sampler::get().untrack_greenlet(greenlet_id);
+
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+link_greenlets(PyObject* Py_UNUSED(m), PyObject* args)
+{
+    uintptr_t parent, child;
+
+    if (!PyArg_ParseTuple(args, "ll", &child, &parent))
+        return NULL;
+
+    Sampler::get().link_greenlets(parent, child);
+
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef _stack_v2_methods[] = {
     { "start", reinterpret_cast<PyCFunction>(stack_v2_start), METH_VARARGS | METH_KEYWORDS, "Start the sampler" },
     { "stop", stack_v2_stop, METH_VARARGS, "Stop the sampler" },
@@ -205,6 +255,11 @@ static PyMethodDef _stack_v2_methods[] = {
     { "track_asyncio_loop", stack_v2_track_asyncio_loop, METH_VARARGS, "Map the name of a task with its identifier" },
     { "init_asyncio", stack_v2_init_asyncio, METH_VARARGS, "Initialise asyncio tracking" },
     { "link_tasks", stack_v2_link_tasks, METH_VARARGS, "Link two tasks" },
+    // greenlet support
+    { "track_greenlet", track_greenlet, METH_VARARGS, "Map a greenlet with its identifier" },
+    { "untrack_greenlet", untrack_greenlet, METH_VARARGS, "Untrack a terminated greenlet" },
+    { "link_greenlets", link_greenlets, METH_VARARGS, "Link two greenlets" },
+
     { "set_adaptive_sampling", stack_v2_set_adaptive_sampling, METH_VARARGS, "Set adaptive sampling" },
     { NULL, NULL, 0, NULL }
 };

ddtrace/profiling/_gevent.py

@@ -0,0 +1,191 @@
+from types import FrameType
+import typing as t
+
+import gevent
+from gevent import thread
+import gevent.greenlet
+from gevent.greenlet import Greenlet as _Greenlet
+import gevent.hub
+from greenlet import greenlet
+from greenlet import settrace
+
+from ddtrace.internal.datadog.profiling import stack_v2
+
+
+# Original objects
+_gevent_hub_spawn_raw = gevent.hub.spawn_raw
+_gevent_joinall = gevent.joinall
+_gevent_wait = gevent.wait
+_gevent_iwait = gevent.iwait
+
+# Global package state
+_greenlet_frame_cells: t.Dict[int, list] = {}
+_original_greenlet_tracer: t.Optional[t.Callable[[str, t.Any], None]] = None
+_greenlet_parent_map: t.Dict[int, int] = {}
+_parent_greenlet_count: t.Dict[int, int] = {}
+
+FRAME_NOT_SET = False  # Sentinel for when the frame is not set
+
+
+def track_gevent_greenlet(greenlet: _Greenlet) -> _Greenlet:
+    greenlet_id = thread.get_ident(greenlet)
+    frame_cell: t.List[t.Union[FrameType, bool, None]] = [FRAME_NOT_SET]
+
+    stack_v2.track_greenlet(  # type: ignore[attr-defined]
+        greenlet_id,
+        greenlet.name or type(greenlet).__qualname__,
+        frame_cell,
+    )
+
+    # Untrack on completion
+    try:
+        greenlet.rawlink(untrack_greenlet)
+    except AttributeError:
+        # This greenlet cannot be linked (e.g. the Hub)
+        pass
+
+    _greenlet_frame_cells[greenlet_id] = frame_cell
+
+    return greenlet
+
+
+def greenlet_tracer(event: str, args: t.Any) -> None:
+    if event in {"switch", "throw"}:
+        # This tracer function runs in the context of the target
+        origin, target = t.cast(tuple, args)
+
+        if (origin_id := thread.get_ident(origin)) not in _greenlet_frame_cells:
+            try:
+                track_gevent_greenlet(origin)
+            except Exception:
+                # Not something that we can track
+                pass
+
+        if (target_id := thread.get_ident(target)) not in _greenlet_frame_cells:
+            # This is likely the hub. We take this chance to track it.
+            try:
+                track_gevent_greenlet(target)
+            except Exception:
+                # Not something that we can track
+                pass
+
+        try:
+            # If this is being set to None, it means the greenlet is likely
+            # finished. We use the sentinel again to signal this.
+            _greenlet_frame_cells[origin_id][0] = origin.gr_frame or FRAME_NOT_SET
+            if target_id not in _parent_greenlet_count:
+                # We don't want to wipe the frame of a parent greenlet because
+                # we need to unwind it. We definitely know it is still running
+                # so if we allow the tracer to set its tracked frame to None,
+                # we won't be able to unwind the full stack.
+                _greenlet_frame_cells[target_id][0] = target.gr_frame  # this *is* None
+        except KeyError:
+            # TODO: Log missing greenlet
+            pass
+
+    if _original_greenlet_tracer is not None:
+        _original_greenlet_tracer(event, args)
+
+
+def untrack_greenlet(greenlet: _Greenlet) -> None:
+    greenlet_id = thread.get_ident(greenlet)
+    stack_v2.untrack_greenlet(greenlet_id)  # type: ignore[attr-defined]
+    _greenlet_frame_cells.pop(greenlet_id, None)
+    _parent_greenlet_count.pop(greenlet_id, None)
+    if (parent_id := _greenlet_parent_map.pop(greenlet_id, None)) is not None:
+        _parent_greenlet_count[parent_id] -= 1
+        if _parent_greenlet_count[parent_id] <= 0:
+            del _parent_greenlet_count[parent_id]
+
+
+def link_greenlets(greenlet_id: int, parent_id: int) -> None:
+    stack_v2.link_greenlets(greenlet_id, parent_id)  # type: ignore[attr-defined]
+    _parent_greenlet_count[parent_id] = _parent_greenlet_count.get(parent_id, 0) + 1
+    _greenlet_parent_map[greenlet_id] = parent_id
+
+
+class Greenlet(_Greenlet):
+    @classmethod
+    def spawn(cls, *args: t.Any, **kwargs: t.Any) -> _Greenlet:
+        return track_gevent_greenlet(super().spawn(*args, **kwargs))
+
+    @classmethod
+    def spawn_later(cls, *args: t.Any, **kwargs: t.Any) -> _Greenlet:
+        return track_gevent_greenlet(super().spawn_later(*args, **kwargs))
+
+    def join(self, *args: t.Any, **kwargs: t.Any) -> None:
+        target_id = thread.get_ident(self)
+        origin_id = thread.get_ident(gevent.getcurrent())
+
+        link_greenlets(target_id, origin_id)
+
+        super().join(*args, **kwargs)
+
+
+def wrap_spawn(original: t.Callable[..., _Greenlet]) -> t.Callable[..., _Greenlet]:
+    def _(*args: t.Any, **kwargs: t.Any) -> _Greenlet:
+        return track_gevent_greenlet(original(*args, **kwargs))
+
+    return _
+
+
+def joinall(greenlets: t.Iterable[_Greenlet], *args: t.Any, **kwargs: t.Any) -> None:
+    # This is a wrapper around gevent.joinall to track the greenlets
+    # that are being joined.
+    current_greenlet = gevent.getcurrent()
+    if isinstance(current_greenlet, greenlet):
+        current_greenlet = gevent.hub.get_hub()
+    current_greenlet_id = thread.get_ident(current_greenlet)
+    for g in greenlets:
+        link_greenlets(thread.get_ident(g), current_greenlet_id)
+    _gevent_joinall(greenlets, *args, **kwargs)
+
+
+def wait_wrapper(original: t.Callable[..., t.Any]) -> t.Callable[..., t.Any]:
+    def _(*args: t.Any, **kwargs: t.Any) -> t.Any:
+        try:
+            objects = args[0]
+        except IndexError:
+            objects = kwargs.get("args", [])
+
+        if greenlets := [_ for _ in objects if isinstance(_, (greenlet, gevent.Greenlet))]:
+            current_greenlet = gevent.getcurrent()
+            if isinstance(current_greenlet, greenlet):
+                current_greenlet = gevent.hub.get_hub()
+            current_greenlet_id = thread.get_ident(current_greenlet)
+            for g in greenlets:
+                link_greenlets(thread.get_ident(g), current_greenlet_id)
+
+        return original(*args, **kwargs)
+
+    return _
+
+
+def patch() -> None:
+    global _original_greenlet_tracer
+
+    # Patch the spawn method to track greenlets.
+    gevent.Greenlet = gevent.greenlet.Greenlet = Greenlet
+    gevent.spawn = Greenlet.spawn
+    gevent.spawn_later = Greenlet.spawn_later
+    gevent.joinall = joinall
+    gevent.wait = wait_wrapper(_gevent_wait)
+    gevent.iwait = wait_wrapper(_gevent_iwait)
+
+    gevent.hub.spawn_raw = wrap_spawn(_gevent_hub_spawn_raw)
+
+    _original_greenlet_tracer = settrace(greenlet_tracer)
+
+
+def unpatch() -> None:
+    # Unpatch the spawn method to stop tracking greenlets.
+    gevent.Greenlet = gevent.greenlet.Greenlet = _Greenlet
+    gevent.spawn = _Greenlet.spawn
+    gevent.spawn_later = _Greenlet.spawn_later
+    gevent.joinall = _gevent_joinall
+    gevent.wait = _gevent_wait
+    gevent.iwait = _gevent_iwait
+
+    gevent.hub.spawn_raw = _gevent_hub_spawn_raw
+
+    settrace(_original_greenlet_tracer)