sampler.hpp

#pragma once

#include <algorithm>
#include <atomic>
#include <condition_variable>
#include <cstdint>
#include <mutex>
#include <random>
#include <vector>

#include "constants.hpp"

#include "echion/task_name.h"
#include "echion/timing.h"

#include <Python.h>

class EchionSampler;

namespace Datadog {

enum class PauseResult : std::uint8_t
{
    Paused,     // sampler was running and is now paused
    NotRunning, // sampler was not running (nothing to pause)
    Timeout,    // sampler is running but did not pause within the timeout
};

class Sampler
{
    // This class manages the initialization of echion as well as the sampling thread.
    // The underlying echion instance it manages keeps much of its state globally, so this class is a singleton in order
    // to keep it aligned with the echion state.
    std::unique_ptr<EchionSampler> echion;

    // The sampling interval is atomic because it needs to be safely propagated to the sampling thread
    std::atomic<microsecond_t> sample_interval_us{ g_default_sampling_period_us };

    // This is not a running total of the number of launched threads; it is a sequence for the
    // transactions upon the sampling threads (usually starts + stops). This allows threads to be
    // stopped or started in a straightforward manner without finer-grained control (locks)
    std::atomic<uint64_t> thread_seq_num{ 0 };

    // Thread exit synchronization - allows stop() to wait for the sampling thread to exit.
    // The mutex + condition variable pair is used to avoid the "lost wake-up" race condition
    // where stop() could miss the notification and hang forever (or until timeout).
    std::atomic<bool> thread_running{ false };
    // Whether the sampler is currently active. Unlike thread_running (which tracks
    // the thread's actual lifecycle) this is set synchronously in start/stop, so
    // it is not subject to the sampling-thread startup race.
    // It becomes false in stop and also when the sampling thread aborts on an
    // unexpected exception, since the sampler is then no longer active.
    // prefork reads this to decide whether to restart the sampler after fork,
    // ensuring a sampler that stopped due to an error is not silently restarted.
    std::atomic<bool> sampler_active_{ false };
    std::mutex thread_exit_mutex;
    std::condition_variable thread_exit_cv;

    // Pause synchronization — allows the faulthandler wrapper to temporarily
    // suspend sampling so signal handlers can be safely swapped without racing
    // with in-flight safe_memcpy calls.
    std::atomic<bool> pause_requested_{ false };
    std::atomic<bool> paused_{ false };
    std::mutex pause_mutex_;
    std::condition_variable pause_cv_;

    // This is a singleton, so no public constructor
    Sampler();

    // One-time setup of echion
    void one_time_setup();

    // Internal perf counters
    uint64_t process_count = 0;
    uint64_t sampler_thread_count = 0;

    bool do_adaptive_sampling = true;
    double target_overhead = g_target_overhead;
    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<PyThreadState> thread_candidates;
    void adapt_sampling_interval();

    // Captures one sampling cycle across all threads (or a reservoir-sampled subset thereof
    // when max_threads_per_sample is set).
    void capture_samples(microsecond_t wall_time_us);

    // Rolling window for p_stable: ring buffer of process_delta values (us CPU per adapt window).
    // p_stable is the p-th percentile of this buffer, giving a stable estimate of app CPU usage
    // that doesn't collapse to zero during brief idle periods.
    std::vector<double> process_delta_window;
    size_t process_delta_window_head = 0;

    // Baseline CPU budget (us per adapt window) corresponding to an absolute floor overhead.
    // Derived from baseline_core_pct at configuration time; keeps the profiler sampling even
    // when app CPU is near 0.
    double baseline_cpu_us_per_adapt_window = 0.0;

    // Percentile (0..1) used for p_stable; configurable, default p95.
    double p_stable_percentile_frac = 0.95;

    // Fast-copy startup warmup in seconds.
    double fast_copy_warmup_seconds = 15.0;

    // Rolling window duration in seconds; controls the ring buffer capacity.
    uint32_t p_stable_window_s = 600;

    // Tracks whether the sampler was running when prefork was called,
    // so that postfork_parent/restart_after_fork can restore it.
    bool was_running_at_fork_{ false };

    void atfork_child();
    friend void stack_atfork_prepare();
    friend void stack_atfork_parent();
    friend void stack_atfork_child();

  public:
    // Singleton instance
    static Sampler& get();

    // Accessor for EchionSampler
    EchionSampler& get_echion() { return *echion; }

    bool start();
    void stop();
    PauseResult pause();
    void resume();
    void register_thread(uint64_t id, uint64_t native_id, const char* name);
    void unregister_thread(uint64_t id);
    void track_asyncio_loop(uintptr_t thread_id, PyObject* loop);
    void init_asyncio(PyObject* _asyncio_scheduled_tasks, PyObject* _asyncio_eager_tasks);
    void link_tasks(PyObject* parent, PyObject* child);
    void weak_link_tasks(PyObject* parent, PyObject* child);
    void sampling_thread(const uint64_t seq_num);
    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 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
    // own intervals accordingly.  Rather than a preemptive measure, we assume the rate is ~fairly stable and just
    // update the next rate with the latest interval. This is not perfect because the adjustment is based on
    // self-time, and we're not currently accounting for the echion self-time.
    void set_interval(double new_interval);
    bool is_running() const { return thread_running.load(); }
    void set_adaptive_sampling(bool value) { do_adaptive_sampling = value; }
    void set_target_overhead(double value) { target_overhead = value; }
    void set_max_sampling_period(microsecond_t max_interval_us)
    {
        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; }

    // Set the absolute overhead floor as "core percent" units (1 = 0.01 core = 10 mcores).
    // Converted to us of CPU budget per adaptation window.
    void set_baseline_core_pct(double value)
    {
        baseline_cpu_us_per_adapt_window = value * 0.01 * g_adaptive_sampling_interval_us;
    }

    // Set the rolling window duration (seconds) over which p_stable is computed.
    void set_p_stable_window_s(uint32_t value) { p_stable_window_s = value; }

    // Set the percentile (0–100) used to compute p_stable from the rolling window.
    void set_p_stable_percentile(double percentile) { p_stable_percentile_frac = percentile / 100.0; }

    void set_fast_copy_warmup_seconds(double value) { fast_copy_warmup_seconds = value; }

    // Delegates to the StackRenderer to clear its caches after fork
    void postfork_child();

    // Record whether the Sampler was running before fork
    void prefork();

    // Restart the sampling thread in the parent after fork
    void postfork_parent();

    // Restart the sampler after fork if it was running.
    // Returns true if start() was invoked and succeeded.
    bool restart_after_fork();
};

void
seed_fast_copy_profiler_stats();

} // namespace Datadog