perf(profiling): consolidate greenlet switch updates

ddtrace/internal/datadog/profiling/stack/__init__.pyi

@@ -98,7 +98,13 @@ def init_asyncio(
 def track_greenlet(greenlet_id: int, name: str, frame: Union[FrameType, bool, None]) -> None: ...
 def untrack_greenlet(greenlet_id: int) -> None: ...
 def link_greenlets(greenlet_id: int, parent_id: int) -> None: ...
-def update_greenlet_frame(greenlet_id: int, frame: Union[FrameType, bool, None]) -> None: ...
+def update_greenlet_switch(
+    origin_id: int,
+    origin_frame: Union[FrameType, bool, None],
+    target_id: int,
+    target_frame: Union[FrameType, bool, None],
+    update_target_frame: bool,
+) -> None: ...
 
 # Module attributes
 is_available: bool

ddtrace/internal/datadog/profiling/stack/_stack.pyi

@@ -66,7 +66,13 @@ def init_asyncio(
 def track_greenlet(greenlet_id: int, name: str, frame: Union[FrameType, bool, None]) -> None: ...
 def untrack_greenlet(greenlet_id: int) -> None: ...
 def link_greenlets(greenlet_id: int, parent_id: int) -> None: ...
-def update_greenlet_frame(greenlet_id: int, frame: Union[FrameType, bool, None]) -> None: ...
+def update_greenlet_switch(
+    origin_id: int,
+    origin_frame: Union[FrameType, bool, None],
+    target_id: int,
+    target_frame: Union[FrameType, bool, None],
+    update_target_frame: bool,
+) -> None: ...
 
 # Native call monitoring (sys.monitoring bridge)
 def start_native_monitoring() -> None: ...

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

@@ -136,7 +136,11 @@ class Sampler
     void track_greenlet(uintptr_t greenlet_id, TaskName name, PyObject* frame);
     void untrack_greenlet(uintptr_t greenlet_id);
     void link_greenlets(uintptr_t parent, uintptr_t child);
-    void update_greenlet_frame(uintptr_t greenlet_id, PyObject* frame);
+    void update_greenlet_switch(uintptr_t origin_id,
+                                PyObject* origin_frame,
+                                uintptr_t target_id,
+                                PyObject* target_frame,
+                                bool update_target_frame);
     void set_uvloop_mode(uintptr_t thread_id, bool value);
 
     // The Python side dynamically adjusts the sampling rate based on overhead, so we need to be able to update our

ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc

@@ -604,7 +604,7 @@ ThreadInfo::unwind_greenlets(EchionSampler& echion, PyThreadState* tstate, unsig
 
     // Phase 1: Snapshot greenlet data under the lock.
     // This minimises the time we hold greenlet_info_map_lock, which is also
-    // acquired by update_greenlet_frame() on every greenlet switch.  Holding
+    // acquired by update_greenlet_switch() on every greenlet switch. Holding
     // the lock during the expensive unwind (Phase 2) would block ALL greenlet
     // switches and lead to resource exhaustion (e.g. DB connection pools).
     {

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

@@ -924,15 +924,22 @@ Sampler::link_greenlets(uintptr_t parent, uintptr_t child)
 }
 
 void
-Sampler::update_greenlet_frame(uintptr_t greenlet_id, PyObject* frame)
+Sampler::update_greenlet_switch(uintptr_t origin_id,
+                                PyObject* origin_frame,
+                                uintptr_t target_id,
+                                PyObject* target_frame,
+                                bool update_target_frame)
 {
     std::lock_guard<std::mutex> guard(echion->greenlet_info_map_lock());
-
     auto& greenlet_info_map = echion->greenlet_info_map();
-    auto entry = greenlet_info_map.find(greenlet_id);
-    if (entry != greenlet_info_map.end()) {
-        // Update the frame of the greenlet
-        entry->second->frame = frame;
+
+    if (auto origin = greenlet_info_map.find(origin_id); origin != greenlet_info_map.end()) {
+        origin->second->frame = origin_frame;
+    }
+    if (update_target_frame) {
+        if (auto target = greenlet_info_map.find(target_id); target != greenlet_info_map.end()) {
+            target->second->frame = target_frame;
+        }
     }
 }
 

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

@@ -510,16 +510,20 @@ link_greenlets(PyObject* Py_UNUSED(m), PyObject* args)
 }
 
 static PyObject*
-update_greenlet_frame(PyObject* Py_UNUSED(m), PyObject* args)
+update_greenlet_switch(PyObject* Py_UNUSED(m), PyObject* args)
 {
-    uintptr_t greenlet_id;
-    PyObject* frame;
+    uintptr_t origin_id;
+    PyObject* origin_frame;
+    uintptr_t target_id;
+    PyObject* target_frame;
+    int update_target_frame;
 
-    if (!PyArg_ParseTuple(args, "lO", &greenlet_id, &frame))
+    if (!PyArg_ParseTuple(args, "lOlOp", &origin_id, &origin_frame, &target_id, &target_frame, &update_target_frame))
         return nullptr;
 
     Py_BEGIN_ALLOW_THREADS;
-    Sampler::get().update_greenlet_frame(greenlet_id, frame);
+    Sampler::get().update_greenlet_switch(
+      origin_id, origin_frame, target_id, target_frame, static_cast<bool>(update_target_frame));
     Py_END_ALLOW_THREADS;
 
     Py_RETURN_NONE;
@@ -1023,7 +1027,7 @@ static PyMethodDef stack_methods[] = {
     { "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" },
-    { "update_greenlet_frame", update_greenlet_frame, METH_VARARGS, "Update the frame of a greenlet" },
+    { "update_greenlet_switch", update_greenlet_switch, METH_VARARGS, "Record a greenlet context switch" },
 
     { "set_adaptive_sampling", stack_set_adaptive_sampling, METH_VARARGS, "Set adaptive sampling" },
     { "set_target_overhead",

ddtrace/profiling/_gevent.py

@@ -66,9 +66,17 @@ def track_gevent_greenlet(gl: _Greenlet, _from_tracer: bool = False) -> _Greenle
     return gl
 
 
-def update_greenlet_frame(greenlet_id: int, frame: t.Union[FrameType, bool, None]) -> None:
-    _tracked_greenlets.add(greenlet_id)
-    stack.update_greenlet_frame(greenlet_id, frame)
+def update_greenlet_switch(
+    origin_id: int,
+    origin_frame: t.Union[FrameType, bool, None],
+    target_id: int,
+    target_frame: t.Union[FrameType, bool, None],
+    update_target_frame: bool,
+) -> None:
+    _tracked_greenlets.add(origin_id)
+    if update_target_frame:
+        _tracked_greenlets.add(target_id)
+    stack.update_greenlet_switch(origin_id, origin_frame, target_id, target_frame, update_target_frame)
 
 
 def greenlet_tracer(event: str, args: t.Any) -> None:
@@ -97,16 +105,18 @@ def greenlet_tracer(event: str, args: t.Any) -> None:
         try:
             # If this is being set to None, it means the greenlet is likely
             # finished. We use the sentinel again to signal this.
-            update_greenlet_frame(
+            origin_frame = t.cast(t.Optional[FrameType], origin.gr_frame) or FRAME_NOT_SET
+            # 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.
+            update_greenlet_switch(
                 origin_id,
-                t.cast(t.Optional[FrameType], origin.gr_frame) or FRAME_NOT_SET,
+                origin_frame,
+                target_id,
+                target.gr_frame,  # This is None for the running target.
+                target_id not in _parent_greenlet_count,
             )
-            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.
-                update_greenlet_frame(target_id, target.gr_frame)  # this *is* None
         except KeyError:
             # TODO: Log missing greenlet
             pass

tests/profiling/collector/test_stack.py

@@ -1099,7 +1099,7 @@ def test_gevent_greenlet_switch_not_blocked_by_profiler() -> None:
 
     Before the fix, unwind_greenlets() held greenlet_info_map_lock for the
     entire stack unwinding of ALL tracked greenlets.  Every greenlet switch
-    calls update_greenlet_frame() under the same lock, so more tracked
+    calls update_greenlet_switch() under the same lock, so more tracked
     greenlets meant longer lock hold and more switch blocking.
 
     This test measures greenlet-switch wall time with zero vs many idle

tests/profiling/test_gevent.py

@@ -11,6 +11,23 @@
 )
 
 
+@pytest.mark.skipif(
+    not GEVENT_COMPATIBLE_WITH_PYTHON_VERSION,
+    reason=f"gevent is not compatible with Python {'.'.join(map(str, tuple(sys.version_info)[:3]))}",
+)
+@pytest.mark.subprocess()
+def test_greenlet_switch_uses_one_native_update() -> None:
+    from unittest.mock import patch
+
+    from ddtrace.profiling import _gevent as _gevent_module
+
+    with patch.object(_gevent_module, "stack") as native_stack:
+        _gevent_module.update_greenlet_switch(1, False, 2, None, True)
+
+    native_stack.update_greenlet_switch.assert_called_once_with(1, False, 2, None, True)
+    assert {1, 2} <= _gevent_module._tracked_greenlets
+
+
 @pytest.mark.skipif(
     not GEVENT_COMPATIBLE_WITH_PYTHON_VERSION,
     reason=f"gevent is not compatible with Python {'.'.join(map(str, tuple(sys.version_info)[:3]))}",