Merge branch 'main' into erikayasuda/test-op

Author: erikayasuda

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

@@ -25,6 +25,7 @@ def set_adaptive_sampling_baseline(baseline_core_pct: float) -> None: ...
 def set_p_stable_window_s(p_stable_window_s: int) -> None: ...
 def set_p_stable_percentile(p_stable_percentile: float) -> None: ...
 def set_max_threads(max_threads: int) -> None: ...
+def set_max_tasks(max_tasks: int) -> None: ...
 def set_uvloop_mode(thread_id: int, uvloop_mode: bool) -> None:
     """Enable uvloop-specific stack unwinding in the native profiler for a specific thread.
 

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

@@ -17,6 +17,7 @@ def set_adaptive_sampling(do_adaptive_sampling: bool = False) -> None: ...
 def set_target_overhead(target_overhead: float) -> None: ...
 def set_max_sampling_period(max_interval_us: int) -> None: ...
 def set_max_threads(max_threads: int) -> None: ...
+def set_max_tasks(max_tasks: int) -> None: ...
 def set_uvloop_mode(thread_id: int, uvloop_mode: bool) -> None: ...
 def set_interval(new_interval: float) -> None: ...
 

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

@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <optional>
+#include <random>
 #include <unordered_map>
 #include <unordered_set>
 
@@ -10,6 +11,7 @@
 #include <echion/strings.h>
 #include <echion/threads.h>
 
+#include "constants.hpp"
 #include "stack_renderer.hpp"
 
 // Forward declaration
@@ -53,6 +55,13 @@ class EchionSampler
     // Only accessed from the sampling thread, so no lock/atomic is needed.
     size_t asyncio_task_count_ = 0;
 
+    // Maximum number of leaf tasks / greenlets to unwind and emit per cycle.
+    // 0 means unlimited.
+    unsigned int max_tasks_per_sample_ = g_default_max_tasks_per_sample;
+
+    // RNG used for task / greenlet reservoir sampling.
+    std::minstd_rand rng_{ std::random_device{}() };
+
     // Caches
     StringTable string_table_;
     LRUCache<uintptr_t, Frame> frame_cache_;
@@ -100,6 +109,11 @@ class EchionSampler
     void add_asyncio_task_count(size_t count) { asyncio_task_count_ += count; }
     size_t asyncio_task_count() const { return asyncio_task_count_; }
 
+    unsigned int max_tasks_per_sample() const { return max_tasks_per_sample_; }
+    void set_max_tasks_per_sample(unsigned int value) { max_tasks_per_sample_ = value; }
+
+    std::minstd_rand& rng() { return rng_; }
+
     // Accessor for StringTable operations
     StringTable& string_table() { return string_table_; }
     const StringTable& string_table() const { return string_table_; }
@@ -138,6 +152,7 @@ class EchionSampler
         asyncio_frame_cache_key_.reset();
         uvloop_frame_cache_key_.reset();
         asyncio_task_count_ = 0;
+        rng_ = std::minstd_rand{ std::random_device{}() };
 
         new (&seen_frames_scratch_) std::unordered_set<PyObject*>();
 

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

@@ -7,6 +7,7 @@
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
 
+#include <optional>
 #include <unordered_set>
 #include <utility>
 #include <vector>
@@ -70,6 +71,10 @@ class StackInfo
     bool on_cpu;
     FrameStack stack;
 
+    // Per-task override wall-time to use in reservoir sampling.
+    // nullopt means "use the thread-level wall time"
+    std::optional<int64_t> walltime_ns = std::nullopt;
+
     StackInfo(TaskName task_name, bool on_cpu, uint64_t task_id)
       : task_name(std::move(task_name))
       , task_id(task_id)

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

@@ -61,7 +61,7 @@ class ThreadInfo
     [[nodiscard]] Result<void> update_cpu_time();
 
     [[nodiscard]] Result<void> sample(EchionSampler&, PyThreadState*, microsecond_t);
-    void unwind(EchionSampler&, PyThreadState*);
+    void unwind(EchionSampler&, PyThreadState*, microsecond_t wall_time_us);
 
     // ------------------------------------------------------------------------
 #if defined PL_LINUX
@@ -116,8 +116,8 @@ class ThreadInfo
   private:
     void reset_cycle_state() noexcept;
     void render_unwound_stacks(EchionSampler&);
-    [[nodiscard]] Result<void> unwind_tasks(EchionSampler&, PyThreadState*);
-    void unwind_greenlets(EchionSampler&, PyThreadState*, unsigned long);
+    [[nodiscard]] Result<void> unwind_tasks(EchionSampler&, PyThreadState*, microsecond_t wall_time_us);
+    void unwind_greenlets(EchionSampler&, PyThreadState*, unsigned long, microsecond_t wall_time_us);
     [[nodiscard]] Result<std::vector<TaskInfo::Ptr>> get_all_tasks(EchionSampler&, PyThreadState* tstate);
 #if PY_VERSION_HEX >= 0x030e0000
     [[nodiscard]] Result<void> get_tasks_from_thread_linked_list(EchionSampler& echion,

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

@@ -37,7 +37,7 @@ LLVMFuzzerTestOneInput(const uint8_t* data, size_t size)
 
     // Only flexes the "unwind thread stack" path, greenlets
     // and asyncio aren't used here
-    thread.unwind(echion_sampler, &tstate);
+    thread.unwind(echion_sampler, &tstate, 0);
 
     g_data = nullptr;
     g_size = 0;

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

@@ -51,7 +51,7 @@ LLVMFuzzerTestOneInput(const uint8_t* data, size_t size)
 
     // Tries to unwind_python_stack (which does nothing because we
     // have bogus data), then unwind_tasks (the one we test here)
-    thread.unwind(echion_sampler, &tstate);
+    thread.unwind(echion_sampler, &tstate, 0);
 
     g_data = nullptr;
     g_size = 0;

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

@@ -15,5 +15,9 @@ constexpr double g_target_overhead = 0.01; // 1% overhead
 // 0 means no limit (sample all threads).
 constexpr unsigned int g_default_max_threads_per_sample = 25;
 
+// Maximum number of leaf tasks / greenlets to sample per cycle. When the number of
+// leaf tasks exceeds this, reservoir sampling is used. 0 means sample all tasks.
+constexpr unsigned int g_default_max_tasks_per_sample = 50;
+
 // Echion maintains a cache of frames--the size of this cache is specified up-front.
 constexpr unsigned int g_default_echion_frame_cache_size = 1024;

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

@@ -156,6 +156,7 @@ class Sampler
         max_sampling_period_us = std::max(max_interval_us, static_cast<microsecond_t>(g_min_sampling_period_us));
     }
     void set_max_threads_per_sample(unsigned int value) { max_threads_per_sample = value; }
+    void set_max_tasks_per_sample(unsigned int value);
 
     // Set the absolute overhead floor as "core percent" units (1 = 0.01 core = 10 mcores).
     // Converted to us of CPU budget per adaptation window.

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

@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <string_view>
 
@@ -87,7 +88,10 @@ class StackRenderer
                              microsecond_t wall_time_us,
                              uintptr_t thread_id,
                              unsigned long native_id);
-    void render_task_begin(std::string_view task_name, bool on_cpu, uint64_t task_id);
+    void render_task_begin(std::string_view task_name,
+                           bool on_cpu,
+                           uint64_t task_id,
+                           std::optional<int64_t> walltime_ns_override = std::nullopt);
     void render_frame(Frame& frame);
     void render_cpu_time(microsecond_t cpu_time_us);
     void render_native_frame(const std::string& name, const std::string& module);