fix(profiling): invalidate stale Echion frames on Python 3.14

Author: taegyunkim

ddtrace/internal/datadog/profiling/stack/echion/echion/echion_sampler.h

@@ -1,9 +1,12 @@
 #pragma once
 
+#include <algorithm>
 #include <cstdint>
 #include <optional>
 #include <unordered_map>
 #include <unordered_set>
+#include <utility>
+#include <vector>
 
 #include <echion/cache.h>
 #include <echion/frame.h>
@@ -56,6 +59,9 @@ class EchionSampler
     // Caches
     StringTable string_table_;
     LRUCache<uintptr_t, Frame> frame_cache_;
+#if PY_VERSION_HEX >= 0x030e0000
+    std::vector<std::pair<int64_t, uint64_t>> code_object_generations_;
+#endif
 
     // Stack renderer for outputting samples
     Datadog::StackRenderer renderer_;
@@ -107,6 +113,46 @@ class EchionSampler
     // Accessor for frame cache operations
     LRUCache<uintptr_t, Frame>& frame_cache() { return frame_cache_; }
 
+    void invalidate_frame_identity_cache()
+    {
+        frame_cache_.clear();
+        asyncio_frame_cache_key_.reset();
+        uvloop_frame_cache_key_.reset();
+    }
+
+#if PY_VERSION_HEX >= 0x030e0000
+    bool update_code_object_generations(const std::vector<InterpreterInfo>& interpreters, bool snapshot_complete)
+    {
+        if (!snapshot_complete || interpreters.empty()) {
+            invalidate_frame_identity_cache();
+            code_object_generations_.clear();
+            return false;
+        }
+
+        bool generations_changed = interpreters.size() != code_object_generations_.size();
+        if (!generations_changed) {
+            for (const auto& interpreter : interpreters) {
+                const std::pair<int64_t, uint64_t> generation{ interpreter.id, interpreter.code_object_generation };
+                if (!std::binary_search(code_object_generations_.begin(), code_object_generations_.end(), generation)) {
+                    generations_changed = true;
+                    break;
+                }
+            }
+        }
+
+        if (generations_changed) {
+            invalidate_frame_identity_cache();
+            code_object_generations_.clear();
+            code_object_generations_.reserve(interpreters.size());
+            for (const auto& interpreter : interpreters) {
+                code_object_generations_.emplace_back(interpreter.id, interpreter.code_object_generation);
+            }
+            std::sort(code_object_generations_.begin(), code_object_generations_.end());
+        }
+        return true;
+    }
+#endif
+
     void postfork_child()
     {
         // Re-init mutexes (placement new to avoid UB)
@@ -121,6 +167,9 @@ class EchionSampler
         // because the Sampling Thread may have been modifying the cache when fork
         // took its snapshot. Traversing a corrupted list to free nodes would crash.
         frame_cache_.postfork_child();
+#if PY_VERSION_HEX >= 0x030e0000
+        new (&code_object_generations_) std::vector<std::pair<int64_t, uint64_t>>();
+#endif
 
         // Also use placement new for all containers touched by the sampling thread.
         // Using placement new means the existing containers are abandoned and

ddtrace/internal/datadog/profiling/stack/echion/echion/interp.h

@@ -27,7 +27,10 @@ class InterpreterInfo
     int64_t id = 0;
     void* tstate_head = NULL;
     void* next = NULL;
+#if PY_VERSION_HEX >= 0x030e0000
+    uint64_t code_object_generation = 0;
+#endif
 };
 
-void
+[[nodiscard]] bool
 for_each_interp(_PyRuntimeState* runtime, const std::function<void(InterpreterInfo& interp)>& callback);

ddtrace/internal/datadog/profiling/stack/fuzz/fuzz_echion_interp.cpp

@@ -28,7 +28,7 @@ LLVMFuzzerTestOneInput(const uint8_t* data, size_t size)
     runtime.interpreters.head = reinterpret_cast<PyInterpreterState*>(p0);
 
     size_t interp_count = 0;
-    for_each_interp(&runtime, [&interp_count](InterpreterInfo&) { interp_count++; });
+    (void)for_each_interp(&runtime, [&interp_count](InterpreterInfo&) { interp_count++; });
 
     g_data = nullptr;
     g_size = 0;

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

@@ -10,6 +10,7 @@
 
 #include "constants.hpp"
 
+#include "echion/interp.h"
 #include "echion/task_name.h"
 #include "echion/timing.h"
 
@@ -79,6 +80,7 @@ class Sampler
     microsecond_t max_sampling_period_us = g_max_sampling_period_us;
     unsigned int max_threads_per_sample = g_default_max_threads_per_sample;
     std::minstd_rand rng{ std::random_device{}() };
+    std::vector<InterpreterInfo> interpreter_candidates;
     std::vector<PyThreadState> thread_candidates;
     void adapt_sampling_interval();
 

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

@@ -1,9 +1,9 @@
 #include <echion/interp.h>
 
-void
+bool
 for_each_interp(_PyRuntimeState* runtime, const std::function<void(InterpreterInfo& interp)>& callback)
 {
-    InterpreterInfo interpreter_info = { 0 };
+    bool snapshot_complete = true;
 
     // Limit interpreter iteration to prevent infinite loops from cycles or corrupted memory.
     // This limit is based on CPython's tachyon profiler (256) and should be more than
@@ -18,15 +18,22 @@ for_each_interp(_PyRuntimeState* runtime, const std::function<void(InterpreterIn
 
         // Cycle detection: if we didn't advance from previous iteration, we're stuck
         if (prev_interp_addr != nullptr && interp_addr == prev_interp_addr) {
-            break; // Cycle detected or failed to advance
+            return false; // Cycle detected or failed to advance
         }
         prev_interp_addr = interp_addr;
 
+        InterpreterInfo interpreter_info = { 0 };
+#if PY_VERSION_HEX >= 0x030e0000
+        snapshot_complete &= !copy_type(interp_addr + offsetof(PyInterpreterState, _code_object_generation),
+                                        interpreter_info.code_object_generation);
+#endif
+
         // Always read next pointer first - we need it to advance
         if (copy_type(interp_addr + offsetof(PyInterpreterState, next), interpreter_info.next))
-            break; // Can't read next, can't advance - stop iteration
+            return false; // Can't read next, can't advance - stop iteration
 
         if (copy_type(interp_addr + offsetof(PyInterpreterState, id), interpreter_info.id)) {
+            snapshot_complete = false;
             interp_addr = reinterpret_cast<char*>(interpreter_info.next);
             continue;
         }
@@ -37,6 +44,7 @@ for_each_interp(_PyRuntimeState* runtime, const std::function<void(InterpreterIn
         if (copy_type(interp_addr + offsetof(PyInterpreterState, tstate_head), interpreter_info.tstate_head))
 #endif
         {
+            snapshot_complete = false;
             interp_addr = reinterpret_cast<char*>(interpreter_info.next);
             continue;
         }
@@ -46,4 +54,6 @@ for_each_interp(_PyRuntimeState* runtime, const std::function<void(InterpreterIn
         // Move to next interpreter
         interp_addr = reinterpret_cast<char*>(interpreter_info.next);
     }
+
+    return snapshot_complete && interp_addr == NULL;
 }

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

@@ -258,26 +258,37 @@ Sampler::capture_samples(const microsecond_t wall_time_us)
 {
     auto* const runtime = &_PyRuntime;
 
+    interpreter_candidates.clear();
+    const bool interpreter_snapshot_complete =
+      for_each_interp(runtime, [&](InterpreterInfo& interp) { interpreter_candidates.push_back(interp); });
+#if PY_VERSION_HEX >= 0x030e0000
+    if (!echion->update_code_object_generations(interpreter_candidates, interpreter_snapshot_complete)) {
+        return;
+    }
+#else
+    (void)interpreter_snapshot_complete;
+#endif
+
     // When max_threads_per_sample is set, we collect all threads first, then apply
     // reservoir sampling (Algorithm R) to select a uniform random subset, and only
     // sample the selected threads. This caps the O(n_threads) stack-unwinding cost.
     if (max_threads_per_sample == 0) {
-        for_each_interp(runtime, [&](InterpreterInfo& interp) -> void {
+        for (auto& interp : interpreter_candidates) {
             for_each_thread(*echion, interp, [&](PyThreadState* tstate, ThreadInfo& thread) {
                 auto success = thread.sample(*echion, tstate, wall_time_us);
                 if (success) {
                     Sample::profile_borrow().stats().increment_sample_count();
                 }
             });
-        });
+        }
     } else {
         thread_candidates.clear();
 
-        for_each_interp(runtime, [&](InterpreterInfo& interp) -> void {
+        for (auto& interp : interpreter_candidates) {
             for_each_thread(*echion, interp, [&](PyThreadState* tstate, ThreadInfo& /*thread*/) {
                 thread_candidates.push_back(*tstate);
             });
-        });
+        }
 
         // Algorithm R: if we have more threads than the cap, select a uniform random subset.
         // Selected threads are placed in [0, sample_count). Overflow threads remain in

ddtrace/internal/datadog/profiling/stack/test/test_sampling_cycle_state.cpp

@@ -5,6 +5,19 @@
 
 #include <memory>
 
+#if PY_VERSION_HEX >= 0x030e0000
+namespace {
+InterpreterInfo
+interpreter(int64_t id, uint64_t generation)
+{
+    InterpreterInfo info;
+    info.id = id;
+    info.code_object_generation = generation;
+    return info;
+}
+} // namespace
+#endif
+
 TEST(SamplingCycleState, UnwindReplacesTaskAndGreenletStacksFromPriorCycle)
 {
     EchionSampler echion;
@@ -23,3 +36,32 @@ TEST(SamplingCycleState, UnwindReplacesTaskAndGreenletStacksFromPriorCycle)
     EXPECT_TRUE(thread.current_tasks.empty());
     EXPECT_TRUE(thread.current_greenlets.empty());
 }
+
+#if PY_VERSION_HEX >= 0x030e0000
+TEST(SamplingCycleState, CodeObjectGenerationInvalidatesFrameIdentityCache)
+{
+    EchionSampler echion(2);
+    constexpr Frame::Key key = 42;
+
+    ASSERT_TRUE(echion.update_code_object_generations({ interpreter(1, 1), interpreter(2, 1) }, true));
+    echion.frame_cache().store(key, std::make_unique<Frame>(10));
+    echion.asyncio_frame_cache_key() = key;
+    echion.uvloop_frame_cache_key() = key;
+
+    EXPECT_TRUE(echion.update_code_object_generations({ interpreter(2, 1), interpreter(1, 1) }, true));
+    EXPECT_TRUE(echion.frame_cache().lookup(key));
+
+    EXPECT_TRUE(echion.update_code_object_generations({ interpreter(1, 1), interpreter(2, 2) }, true));
+    EXPECT_FALSE(echion.frame_cache().lookup(key));
+    EXPECT_FALSE(echion.asyncio_frame_cache_key());
+    EXPECT_FALSE(echion.uvloop_frame_cache_key());
+
+    echion.frame_cache().store(key, std::make_unique<Frame>(10));
+    EXPECT_TRUE(echion.update_code_object_generations({ interpreter(1, 1), interpreter(3, 2) }, true));
+    EXPECT_FALSE(echion.frame_cache().lookup(key));
+
+    echion.frame_cache().store(key, std::make_unique<Frame>(10));
+    EXPECT_FALSE(echion.update_code_object_generations({ interpreter(1, 1), interpreter(3, 2) }, false));
+    EXPECT_FALSE(echion.frame_cache().lookup(key));
+}
+#endif

releasenotes/notes/profiling-echion-code-generation-6df0c11c31a5d70c.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    profiling: On Python 3.14, prevents stack samples from being attributed to stale Python frames after code objects are replaced.

tests/profiling/collector/test_stack.py

@@ -150,6 +150,77 @@ def foo() -> None:
     pprof_utils.assert_profile_has_sample(profile, samples=samples, expected_sample=expected_sample)
 
 
+@pytest.mark.skipif(sys.version_info < (3, 14), reason="requires CPython's code object generation")
+@pytest.mark.subprocess()
+def test_code_object_address_reuse_does_not_return_stale_frame() -> None:
+    import gc
+    import os
+    from pathlib import Path
+    import tempfile
+    import time
+    from types import FunctionType
+    import weakref
+
+    import _interpreters
+
+    from ddtrace.internal.datadog.profiling import ddup
+    from ddtrace.profiling.collector import stack
+    from tests.profiling.collector import pprof_utils
+
+    interpreter = _interpreters.create("legacy")
+    test_name = "test_code_object_address_reuse_does_not_return_stale_frame"
+    tmp_path = Path(tempfile.mkdtemp(prefix=test_name))
+    pprof_prefix = str(tmp_path / test_name)
+    output_filename = pprof_prefix + "." + str(os.getpid())
+    code_filename = "echion-code-reuse.py"
+
+    assert ddup.is_available
+    ddup.config(env="test", service=test_name, version="my_version", output_filename=pprof_prefix)
+    ddup.start()
+    ddup.upload()
+
+    namespace = {"time": time}
+    source = "def template(deadline):\n    while time.monotonic() < deadline:\n        pass\n"
+    exec(compile(source, code_filename, "exec"), namespace)
+    template_code = namespace["template"].__code__
+
+    def make_function(name: str):
+        code = template_code.replace(co_name=name, co_qualname=name)
+        return FunctionType(code, {"time": time})
+
+    old_name = "old_dynamic_function"
+    old_function = make_function(old_name)
+
+    with stack.StackCollector():
+        old_function(time.monotonic() + 0.3)
+        ddup.upload()
+
+        old_address = id(old_function.__code__)
+        old_code = weakref.ref(old_function.__code__)
+        del old_function
+        gc.collect()
+        assert old_code() is None
+
+        replacement_name = "new_dynamic_function"
+        replacement_function = make_function(replacement_name)
+        assert id(replacement_function.__code__) == old_address
+        replacement_function(time.monotonic() + 0.3)
+
+    ddup.upload()
+
+    profile = pprof_utils.parse_newest_profile(output_filename)
+    samples = pprof_utils.get_samples_with_value_type(profile, "wall-time")
+    sampled_names = {
+        location.function_name
+        for sample in samples
+        for location in (pprof_utils.get_location_from_id(profile, location_id) for location_id in sample.location_id)
+        if location.filename == code_filename
+    }
+    assert replacement_name in sampled_names
+    assert old_name not in sampled_names
+    _interpreters.destroy(interpreter)
+
+
 def test_push_span(tmp_path: Path, tracer: Tracer) -> None:
     test_name = "test_push_span"
     pprof_prefix = str(tmp_path / test_name)