diff --git a/ddtrace/internal/datadog/profiling/stack/__init__.pyi b/ddtrace/internal/datadog/profiling/stack/__init__.pyi index b41e62760f2..134657dfb29 100644 --- a/ddtrace/internal/datadog/profiling/stack/__init__.pyi +++ b/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 record_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 diff --git a/ddtrace/internal/datadog/profiling/stack/_stack.pyi b/ddtrace/internal/datadog/profiling/stack/_stack.pyi index b5abfecba4a..9b976d01439 100644 --- a/ddtrace/internal/datadog/profiling/stack/_stack.pyi +++ b/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 record_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: ... diff --git a/ddtrace/internal/datadog/profiling/stack/include/sampler.hpp b/ddtrace/internal/datadog/profiling/stack/include/sampler.hpp index 9d3b4490c07..363a4257d15 100644 --- a/ddtrace/internal/datadog/profiling/stack/include/sampler.hpp +++ b/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 record_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 diff --git a/ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc b/ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc index 980f466dcb1..18c13599ffb 100644 --- a/ddtrace/internal/datadog/profiling/stack/src/echion/threads.cc +++ b/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 record_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). { diff --git a/ddtrace/internal/datadog/profiling/stack/src/sampler.cpp b/ddtrace/internal/datadog/profiling/stack/src/sampler.cpp index c0f36cc2313..d7df2e5574c 100644 --- a/ddtrace/internal/datadog/profiling/stack/src/sampler.cpp +++ b/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::record_greenlet_switch(uintptr_t origin_id, + PyObject* origin_frame, + uintptr_t target_id, + PyObject* target_frame, + bool update_target_frame) { std::lock_guard 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; + } } } diff --git a/ddtrace/internal/datadog/profiling/stack/src/stack.cpp b/ddtrace/internal/datadog/profiling/stack/src/stack.cpp index 45b1347f37b..68b17abbf83 100644 --- a/ddtrace/internal/datadog/profiling/stack/src/stack.cpp +++ b/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) +record_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().record_greenlet_switch( + origin_id, origin_frame, target_id, target_frame, static_cast(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" }, + { "record_greenlet_switch", record_greenlet_switch, METH_VARARGS, "Record a greenlet context switch" }, { "set_adaptive_sampling", stack_set_adaptive_sampling, METH_VARARGS, "Set adaptive sampling" }, { "set_target_overhead", diff --git a/ddtrace/internal/datadog/profiling/stack/test/test_sampling_cycle_state.cpp b/ddtrace/internal/datadog/profiling/stack/test/test_sampling_cycle_state.cpp index 8fa3bad7556..5d0419189d1 100644 --- a/ddtrace/internal/datadog/profiling/stack/test/test_sampling_cycle_state.cpp +++ b/ddtrace/internal/datadog/profiling/stack/test/test_sampling_cycle_state.cpp @@ -1,5 +1,6 @@ #include "echion/echion_sampler.h" #include "echion/task_name.h" +#include "sampler.hpp" #include @@ -23,3 +24,44 @@ TEST(SamplingCycleState, UnwindReplacesTaskAndGreenletStacksFromPriorCycle) EXPECT_TRUE(thread.current_tasks.empty()); EXPECT_TRUE(thread.current_greenlets.empty()); } + +TEST(SamplingCycleState, GreenletSwitchPreservesLinkedParentFrame) +{ + constexpr GreenletInfo::ID child_id = 101; + constexpr GreenletInfo::ID parent_id = 102; + PyObject child_running_frame{}; + PyObject child_suspended_frame{}; + PyObject parent_suspended_frame{}; + PyObject parent_resumed_frame{}; + + Datadog::Sampler& sampler = Datadog::Sampler::get(); + EchionSampler& echion = sampler.get_echion(); + { + std::lock_guard guard(echion.greenlet_info_map_lock()); + auto& greenlets = echion.greenlet_info_map(); + greenlets.emplace( + child_id, std::make_unique(child_id, &child_running_frame, TaskName::from_literal("child"))); + greenlets.emplace( + parent_id, + std::make_unique(parent_id, &parent_suspended_frame, TaskName::from_literal("parent"))); + } + sampler.link_greenlets(parent_id, child_id); + + sampler.record_greenlet_switch(child_id, &child_suspended_frame, parent_id, &parent_resumed_frame, false); + { + std::lock_guard guard(echion.greenlet_info_map_lock()); + EXPECT_EQ(echion.greenlet_parent_map().at(child_id), parent_id); + EXPECT_EQ(echion.greenlet_info_map().at(child_id)->frame, &child_suspended_frame); + EXPECT_EQ(echion.greenlet_info_map().at(parent_id)->frame, &parent_suspended_frame); + } + + sampler.record_greenlet_switch(child_id, &child_running_frame, parent_id, &parent_resumed_frame, true); + { + std::lock_guard guard(echion.greenlet_info_map_lock()); + EXPECT_EQ(echion.greenlet_info_map().at(child_id)->frame, &child_running_frame); + EXPECT_EQ(echion.greenlet_info_map().at(parent_id)->frame, &parent_resumed_frame); + echion.greenlet_info_map().erase(child_id); + echion.greenlet_info_map().erase(parent_id); + echion.greenlet_parent_map().erase(child_id); + } +} diff --git a/ddtrace/profiling/_gevent.py b/ddtrace/profiling/_gevent.py index d943ac186f0..37dd98ceeef 100644 --- a/ddtrace/profiling/_gevent.py +++ b/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 record_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.record_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. + record_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 diff --git a/tests/profiling/collector/test_stack.py b/tests/profiling/collector/test_stack.py index 065af199827..ef7c35d7f65 100644 --- a/tests/profiling/collector/test_stack.py +++ b/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 record_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