ddtrace 1.20.0 → 1.23.0

data/ext/{ddtrace_profiling_native_extension → datadog_profiling_native_extension}/collectors_cpu_and_wall_time_worker.c

@@ -17,6 +17,8 @@
 #include "setup_signal_handler.h"
 #include "time_helpers.h"
 
+#define ERR_CLOCK_FAIL "failed to get clock time"
+
 // Maximum allowed value for an allocation weight. Attempts to use higher values will result in clamping.
 unsigned int MAX_ALLOC_WEIGHT = 65535;
 
@@ -94,6 +96,7 @@ struct cpu_and_wall_time_worker_state {
   bool no_signals_workaround_enabled;
   bool dynamic_sampling_rate_enabled;
   bool allocation_profiling_enabled;
+  bool skip_idle_samples_for_testing;
   VALUE self_instance;
   VALUE thread_context_collector_instance;
   VALUE idle_sampling_helper_instance;
@@ -119,6 +122,7 @@ struct cpu_and_wall_time_worker_state {
   bool during_sample;
 
   struct stats {
+    // # Generic stats
     // How many times we tried to trigger a sample
     unsigned int trigger_sample_attempts;
     // How many times we tried to simulate signal delivery
@@ -129,25 +133,38 @@ struct cpu_and_wall_time_worker_state {
     unsigned int signal_handler_enqueued_sample;
     // How many times the signal handler was called from the wrong thread
     unsigned int signal_handler_wrong_thread;
-    // How many times we actually sampled (except GC samples)
-    unsigned int sampled;
-    // How many times we skipped a sample because of the dynamic sampling rate mechanism
-    unsigned int skipped_sample_because_of_dynamic_sampling_rate;
+    // How many times we actually tried to interrupt a thread for sampling
+    unsigned int interrupt_thread_attempts;
 
-    // Stats for the results of calling rb_postponed_job_register_one
-      // The same function was already waiting to be executed
+    // # Stats for the results of calling rb_postponed_job_register_one
+    // The same function was already waiting to be executed
     unsigned int postponed_job_skipped_already_existed;
-      // The function was added to the queue successfully
+    // The function was added to the queue successfully
     unsigned int postponed_job_success;
-      // The queue was full
+    // The queue was full
     unsigned int postponed_job_full;
-      // The function returned an unknown result code
+    // The function returned an unknown result code
     unsigned int postponed_job_unknown_result;
 
-    // Min/max/total wall-time spent sampling (except GC samples)
-    uint64_t sampling_time_ns_min;
-    uint64_t sampling_time_ns_max;
-    uint64_t sampling_time_ns_total;
+    // # CPU/Walltime sampling stats
+    // How many times we actually CPU/wall sampled
+    unsigned int cpu_sampled;
+    // How many times we skipped a CPU/wall sample because of the dynamic sampling rate mechanism
+    unsigned int cpu_skipped;
+    // Min/max/total wall-time spent on CPU/wall sampling
+    uint64_t cpu_sampling_time_ns_min;
+    uint64_t cpu_sampling_time_ns_max;
+    uint64_t cpu_sampling_time_ns_total;
+
+    // # Allocation sampling stats
+    // How many times we actually allocation sampled
+    uint64_t allocation_sampled;
+    // How many times we skipped an allocation sample because of the dynamic sampling rate mechanism
+    uint64_t allocation_skipped;
+    // Min/max/total wall-time spent on allocation sampling
+    uint64_t allocation_sampling_time_ns_min;
+    uint64_t allocation_sampling_time_ns_max;
+    uint64_t allocation_sampling_time_ns_total;
     // How many times we saw allocations being done inside a sample
     unsigned int allocations_during_sample;
   } stats;
@@ -163,12 +180,14 @@ static VALUE _native_initialize(
   VALUE no_signals_workaround_enabled,
   VALUE dynamic_sampling_rate_enabled,
   VALUE dynamic_sampling_rate_overhead_target_percentage,
-  VALUE allocation_profiling_enabled
+  VALUE allocation_profiling_enabled,
+  VALUE skip_idle_samples_for_testing
 );
 static void cpu_and_wall_time_worker_typed_data_mark(void *state_ptr);
 static VALUE _native_sampling_loop(VALUE self, VALUE instance);
 static VALUE _native_stop(DDTRACE_UNUSED VALUE _self, VALUE self_instance, VALUE worker_thread);
 static VALUE stop(VALUE self_instance, VALUE optional_exception);
+static void stop_state(struct cpu_and_wall_time_worker_state *state, VALUE optional_exception);
 static void handle_sampling_signal(DDTRACE_UNUSED int _signal, DDTRACE_UNUSED siginfo_t *_info, DDTRACE_UNUSED void *_ucontext);
 static void *run_sampling_trigger_loop(void *state_ptr);
 static void interrupt_sampling_trigger_loop(void *state_ptr);
@@ -191,15 +210,18 @@ static VALUE _native_simulate_sample_from_postponed_job(DDTRACE_UNUSED VALUE sel
 static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE instance);
 static VALUE _native_is_sigprof_blocked_in_current_thread(DDTRACE_UNUSED VALUE self);
 static VALUE _native_stats(DDTRACE_UNUSED VALUE self, VALUE instance);
+static VALUE _native_stats_reset_not_thread_safe(DDTRACE_UNUSED VALUE self, VALUE instance);
 void *simulate_sampling_signal_delivery(DDTRACE_UNUSED void *_unused);
 static void grab_gvl_and_sample(void);
-static void reset_stats(struct cpu_and_wall_time_worker_state *state);
+static void reset_stats_not_thread_safe(struct cpu_and_wall_time_worker_state *state);
 static void sleep_for(uint64_t time_ns);
 static VALUE _native_allocation_count(DDTRACE_UNUSED VALUE self);
 static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused);
 static void disable_tracepoints(struct cpu_and_wall_time_worker_state *state);
 static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self);
 static VALUE rescued_sample_allocation(VALUE tracepoint_data);
+static void delayed_error(struct cpu_and_wall_time_worker_state *state, const char *error);
+static VALUE _native_delayed_error(DDTRACE_UNUSED VALUE self, VALUE instance, VALUE error_msg);
 
 // Note on sampler global state safety:
 //
@@ -212,6 +234,11 @@ static VALUE rescued_sample_allocation(VALUE tracepoint_data);
 static VALUE active_sampler_instance = Qnil;
 static struct cpu_and_wall_time_worker_state *active_sampler_instance_state = NULL;
 
+// See handle_sampling_signal for details on what this does
+#ifdef NO_POSTPONED_TRIGGER
+  static void *gc_finalize_deferred_workaround;
+#endif
+
 // Used to implement CpuAndWallTimeWorker._native_allocation_count . To be able to use cheap thread-local variables
 // (here with `__thread`, see https://gcc.gnu.org/onlinedocs/gcc/Thread-Local.html), this needs to be global.
 //
@@ -230,6 +257,8 @@ void collectors_cpu_and_wall_time_worker_init(VALUE profiling_module) {
     if (sample_from_postponed_job_handle == POSTPONED_JOB_HANDLE_INVALID || after_gc_from_postponed_job_handle == POSTPONED_JOB_HANDLE_INVALID) {
       rb_raise(rb_eRuntimeError, "Failed to register profiler postponed jobs (got POSTPONED_JOB_HANDLE_INVALID)");
     }
+  #else
+    gc_finalize_deferred_workaround = objspace_ptr_for_gc_finalize_deferred_workaround();
   #endif
 
   VALUE collectors_module = rb_define_module_under(profiling_module, "Collectors");
@@ -247,13 +276,16 @@ void collectors_cpu_and_wall_time_worker_init(VALUE profiling_module) {
   // https://bugs.ruby-lang.org/issues/18007 for a discussion around this.
   rb_define_alloc_func(collectors_cpu_and_wall_time_worker_class, _native_new);
 
-  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_initialize", _native_initialize, 8);
+  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_initialize", _native_initialize, 9);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_sampling_loop", _native_sampling_loop, 1);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_stop", _native_stop, 2);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_reset_after_fork", _native_reset_after_fork, 1);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_stats", _native_stats, 1);
+  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_stats_reset_not_thread_safe", _native_stats_reset_not_thread_safe, 1);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_allocation_count", _native_allocation_count, 0);
+  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_is_running?", _native_is_running, 1);
   rb_define_singleton_method(testing_module, "_native_current_sigprof_signal_handler", _native_current_sigprof_signal_handler, 0);
+  // TODO: Remove `_native_is_running` from `testing_module` once `prof-correctness` has been updated to not need it
   rb_define_singleton_method(testing_module, "_native_is_running?", _native_is_running, 1);
   rb_define_singleton_method(testing_module, "_native_install_testing_signal_handler", _native_install_testing_signal_handler, 0);
   rb_define_singleton_method(testing_module, "_native_remove_testing_signal_handler", _native_remove_testing_signal_handler, 0);
@@ -263,6 +295,7 @@ void collectors_cpu_and_wall_time_worker_init(VALUE profiling_module) {
   rb_define_singleton_method(testing_module, "_native_simulate_sample_from_postponed_job", _native_simulate_sample_from_postponed_job, 0);
   rb_define_singleton_method(testing_module, "_native_is_sigprof_blocked_in_current_thread", _native_is_sigprof_blocked_in_current_thread, 0);
   rb_define_singleton_method(testing_module, "_native_with_blocked_sigprof", _native_with_blocked_sigprof, 0);
+  rb_define_singleton_method(testing_module, "_native_delayed_error", _native_delayed_error, 2);
 }
 
 // This structure is used to define a Ruby object that stores a pointer to a struct cpu_and_wall_time_worker_state
@@ -288,11 +321,11 @@ static VALUE _native_new(VALUE klass) {
   state->no_signals_workaround_enabled = false;
   state->dynamic_sampling_rate_enabled = true;
   state->allocation_profiling_enabled = false;
+  state->skip_idle_samples_for_testing = false;
   state->thread_context_collector_instance = Qnil;
   state->idle_sampling_helper_instance = Qnil;
   state->owner_thread = Qnil;
   dynamic_sampling_rate_init(&state->cpu_dynamic_sampling_rate);
-  discrete_dynamic_sampler_init(&state->allocation_sampler, "allocation");
   state->gc_tracepoint = Qnil;
   state->object_allocation_tracepoint = Qnil;
 
@@ -302,7 +335,15 @@ static VALUE _native_new(VALUE klass) {
 
   state->during_sample = false;
 
-  reset_stats(state);
+  reset_stats_not_thread_safe(state);
+
+  long now = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
+  if (now == 0) {
+    ruby_xfree(state);
+    rb_raise(rb_eRuntimeError, ERR_CLOCK_FAIL);
+  }
+
+  discrete_dynamic_sampler_init(&state->allocation_sampler, "allocation", now);
 
   return state->self_instance = TypedData_Wrap_Struct(klass, &cpu_and_wall_time_worker_typed_data, state);
 }
@@ -316,13 +357,15 @@ static VALUE _native_initialize(
   VALUE no_signals_workaround_enabled,
   VALUE dynamic_sampling_rate_enabled,
   VALUE dynamic_sampling_rate_overhead_target_percentage,
-  VALUE allocation_profiling_enabled
+  VALUE allocation_profiling_enabled,
+  VALUE skip_idle_samples_for_testing
 ) {
   ENFORCE_BOOLEAN(gc_profiling_enabled);
   ENFORCE_BOOLEAN(no_signals_workaround_enabled);
   ENFORCE_BOOLEAN(dynamic_sampling_rate_enabled);
   ENFORCE_TYPE(dynamic_sampling_rate_overhead_target_percentage, T_FLOAT);
   ENFORCE_BOOLEAN(allocation_profiling_enabled);
+  ENFORCE_BOOLEAN(skip_idle_samples_for_testing)
 
   struct cpu_and_wall_time_worker_state *state;
   TypedData_Get_Struct(self_instance, struct cpu_and_wall_time_worker_state, &cpu_and_wall_time_worker_typed_data, state);
@@ -331,6 +374,7 @@ static VALUE _native_initialize(
   state->no_signals_workaround_enabled = (no_signals_workaround_enabled == Qtrue);
   state->dynamic_sampling_rate_enabled = (dynamic_sampling_rate_enabled == Qtrue);
   state->allocation_profiling_enabled = (allocation_profiling_enabled == Qtrue);
+  state->skip_idle_samples_for_testing = (skip_idle_samples_for_testing == Qtrue);
 
   double total_overhead_target_percentage = NUM2DBL(dynamic_sampling_rate_overhead_target_percentage);
   if (!state->allocation_profiling_enabled) {
@@ -339,7 +383,8 @@ static VALUE _native_initialize(
     // TODO: May be nice to offer customization here? Distribute available "overhead" margin with a bias towards one or the other
     // sampler.
     dynamic_sampling_rate_set_overhead_target_percentage(&state->cpu_dynamic_sampling_rate, total_overhead_target_percentage / 2);
-    discrete_dynamic_sampler_set_overhead_target_percentage(&state->allocation_sampler, total_overhead_target_percentage / 2);
+    long now = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
+    discrete_dynamic_sampler_set_overhead_target_percentage(&state->allocation_sampler, total_overhead_target_percentage / 2, now);
   }
 
   state->thread_context_collector_instance = enforce_thread_context_collector_instance(thread_context_collector_instance);
@@ -368,6 +413,12 @@ static VALUE _native_sampling_loop(DDTRACE_UNUSED VALUE _self, VALUE instance) {
   struct cpu_and_wall_time_worker_state *state;
   TypedData_Get_Struct(instance, struct cpu_and_wall_time_worker_state, &cpu_and_wall_time_worker_typed_data, state);
 
+  // If we already got a delayed exception registered even before starting, raise before starting
+  if (state->failure_exception != Qnil) {
+    disable_tracepoints(state);
+    rb_exc_raise(state->failure_exception);
+  }
+
   struct cpu_and_wall_time_worker_state *old_state = active_sampler_instance_state;
   if (old_state != NULL) {
     if (is_thread_alive(old_state->owner_thread)) {
@@ -393,7 +444,8 @@ static VALUE _native_sampling_loop(DDTRACE_UNUSED VALUE _self, VALUE instance) {
 
   // Reset the dynamic sampling rate state, if any (reminder: the monotonic clock reference may change after a fork)
   dynamic_sampling_rate_reset(&state->cpu_dynamic_sampling_rate);
-  discrete_dynamic_sampler_reset(&state->allocation_sampler);
+  long now = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
+  discrete_dynamic_sampler_reset(&state->allocation_sampler, now);
 
   // This write to a global is thread-safe BECAUSE we're still holding on to the global VM lock at this point
   active_sampler_instance_state = state;
@@ -455,15 +507,19 @@ static VALUE _native_stop(DDTRACE_UNUSED VALUE _self, VALUE self_instance, VALUE
   return stop(self_instance, /* optional_exception: */ Qnil);
 }
 
-static VALUE stop(VALUE self_instance, VALUE optional_exception) {
-  struct cpu_and_wall_time_worker_state *state;
-  TypedData_Get_Struct(self_instance, struct cpu_and_wall_time_worker_state, &cpu_and_wall_time_worker_typed_data, state);
-
+static void stop_state(struct cpu_and_wall_time_worker_state *state, VALUE optional_exception) {
   atomic_store(&state->should_run, false);
   state->failure_exception = optional_exception;
 
   // Disable the tracepoints as soon as possible, so the VM doesn't keep on calling them
   disable_tracepoints(state);
+}
+
+static VALUE stop(VALUE self_instance, VALUE optional_exception) {
+  struct cpu_and_wall_time_worker_state *state;
+  TypedData_Get_Struct(self_instance, struct cpu_and_wall_time_worker_state, &cpu_and_wall_time_worker_typed_data, state);
+
+  stop_state(state, optional_exception);
 
   return Qtrue;
 }
@@ -503,7 +559,32 @@ static void handle_sampling_signal(DDTRACE_UNUSED int _signal, DDTRACE_UNUSED si
     rb_postponed_job_trigger(sample_from_postponed_job_handle);
     state->stats.postponed_job_success++; // Always succeeds
   #else
-    int result = rb_postponed_job_register_one(0, sample_from_postponed_job, NULL);
+
+    // This is a workaround for https://bugs.ruby-lang.org/issues/19991 (for Ruby < 3.3)
+    //
+    // TL;DR the `rb_postponed_job_register_one` API is not atomic (which is why it got replaced by `rb_postponed_job_trigger`)
+    // and in rare cases can cause VM crashes.
+    //
+    // Specifically, if we're interrupting `rb_postponed_job_flush` (the function that processes postponed jobs), the way
+    // that this function reads the jobs is not atomic, and can cause our call to
+    // `rb_postponed_job_register(function, arg)` to clobber an existing job that is getting dequeued.
+    // Clobbering an existing job is somewhat annoying, but the worst part is that it can happen that we clobber only
+    // the existing job's arguments.
+    // As surveyed in https://github.com/ruby/ruby/pull/8949#issuecomment-1821441370 clobbering the arguments turns out
+    // to not matter in many cases as usually `rb_postponed_job_register` calls in the VM and ecosystem ignore the argument.
+    //
+    // https://bugs.ruby-lang.org/issues/19991 is the exception: inside Ruby's `gc.c`, when dealing with object
+    // finalizers, Ruby calls `gc_finalize_deferred_register` which internally calls
+    // `rb_postponed_job_register_one(gc_finalize_deferred, objspace)`.
+    // Clobbering this call means that `gc_finalize_deferred` would get called with `NULL`, causing a segmentation fault.
+    //
+    // Note that this is quite rare: our signal needs to land at exactly the point where the VM has read the function
+    // to execute, but has yet to read the arguments. @ivoanjo: I could only reproduce it by manually changing the VM
+    // code to simulate this happening.
+    //
+    // Thus, our workaround is simple: we pass in objspace as our argument, just in case the clobbering happens.
+    // In the happy path, we never use this argument so it makes no difference. In the buggy path, we avoid crashing the VM.
+    int result = rb_postponed_job_register(0, sample_from_postponed_job, gc_finalize_deferred_workaround /* instead of NULL */);
 
     // Officially, the result of rb_postponed_job_register_one is documented as being opaque, but in practice it does not
     // seem to have changed between Ruby 2.3 and 3.2, and so we track it as a debugging mechanism
@@ -545,17 +626,23 @@ static void *run_sampling_trigger_loop(void *state_ptr) {
         // Note that reading the GVL owner and sending them a signal is a race -- the Ruby VM keeps on executing while
         // we're doing this, so we may still not signal the correct thread from time to time, but our signal handler
         // includes a check to see if it got called in the right thread
+        state->stats.interrupt_thread_attempts++;
         pthread_kill(owner.owner, SIGPROF);
       } else {
-        // If no thread owns the Global VM Lock, the application is probably idle at the moment. We still want to sample
-        // so we "ask a friend" (the IdleSamplingHelper component) to grab the GVL and simulate getting a SIGPROF.
-        //
-        // In a previous version of the code, we called `grab_gvl_and_sample` directly BUT this was problematic because
-        // Ruby may concurrently get busy and so the CpuAndWallTimeWorker would be blocked in line to acquire the GVL
-        // for an uncontrolled amount of time. (This can still happen to the IdleSamplingHelper, but the
-        // CpuAndWallTimeWorker will still be free to interrupt the Ruby VM and keep sampling for the entire blocking period).
-        state->stats.trigger_simulated_signal_delivery_attempts++;
-        idle_sampling_helper_request_action(state->idle_sampling_helper_instance, grab_gvl_and_sample);
+        if (state->skip_idle_samples_for_testing) {
+          // This was added to make sure our tests don't accidentally pass due to idle samples. Specifically, if we
+          // comment out the thread interruption code inside `if (owner.valid)` above, our tests should not pass!
+        } else {
+          // If no thread owns the Global VM Lock, the application is probably idle at the moment. We still want to sample
+          // so we "ask a friend" (the IdleSamplingHelper component) to grab the GVL and simulate getting a SIGPROF.
+          //
+          // In a previous version of the code, we called `grab_gvl_and_sample` directly BUT this was problematic because
+          // Ruby may concurrently get busy and so the CpuAndWallTimeWorker would be blocked in line to acquire the GVL
+          // for an uncontrolled amount of time. (This can still happen to the IdleSamplingHelper, but the
+          // CpuAndWallTimeWorker will still be free to interrupt the Ruby VM and keep sampling for the entire blocking period).
+          state->stats.trigger_simulated_signal_delivery_attempts++;
+          idle_sampling_helper_request_action(state->idle_sampling_helper_instance, grab_gvl_and_sample);
+        }
       }
     }
 
@@ -607,11 +694,11 @@ static VALUE rescued_sample_from_postponed_job(VALUE self_instance) {
   long wall_time_ns_before_sample = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
 
   if (state->dynamic_sampling_rate_enabled && !dynamic_sampling_rate_should_sample(&state->cpu_dynamic_sampling_rate, wall_time_ns_before_sample)) {
-    state->stats.skipped_sample_because_of_dynamic_sampling_rate++;
+    state->stats.cpu_skipped++;
     return Qnil;
   }
 
-  state->stats.sampled++;
+  state->stats.cpu_sampled++;
 
   VALUE profiler_overhead_stack_thread = state->owner_thread; // Used to attribute profiler overhead to a different stack
   thread_context_collector_sample(state->thread_context_collector_instance, wall_time_ns_before_sample, profiler_overhead_stack_thread);
@@ -622,9 +709,9 @@ static VALUE rescued_sample_from_postponed_job(VALUE self_instance) {
   // Guard against wall-time going backwards, see https://github.com/DataDog/dd-trace-rb/pull/2336 for discussion.
   uint64_t sampling_time_ns = delta_ns < 0 ? 0 : delta_ns;
 
-  state->stats.sampling_time_ns_min = uint64_min_of(sampling_time_ns, state->stats.sampling_time_ns_min);
-  state->stats.sampling_time_ns_max = uint64_max_of(sampling_time_ns, state->stats.sampling_time_ns_max);
-  state->stats.sampling_time_ns_total += sampling_time_ns;
+  state->stats.cpu_sampling_time_ns_min = uint64_min_of(sampling_time_ns, state->stats.cpu_sampling_time_ns_min);
+  state->stats.cpu_sampling_time_ns_max = uint64_max_of(sampling_time_ns, state->stats.cpu_sampling_time_ns_max);
+  state->stats.cpu_sampling_time_ns_total += sampling_time_ns;
 
   dynamic_sampling_rate_after_sample(&state->cpu_dynamic_sampling_rate, wall_time_ns_after_sample, sampling_time_ns);
 
@@ -666,6 +753,9 @@ static VALUE release_gvl_and_run_sampling_trigger_loop(VALUE instance) {
   if (state->gc_profiling_enabled) rb_tracepoint_enable(state->gc_tracepoint);
   if (state->allocation_profiling_enabled) rb_tracepoint_enable(state->object_allocation_tracepoint);
 
+  // Flag the profiler as running before we release the GVL, in case anyone's waiting to know about it
+  rb_funcall(instance, rb_intern("signal_running"), 0);
+
   rb_thread_call_without_gvl(run_sampling_trigger_loop, state, interrupt_sampling_trigger_loop, state);
 
   // If we stopped sampling due to an exception, re-raise it (now in the worker thread)
@@ -823,7 +913,7 @@ static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE instance)
   // Disable all tracepoints, so that there are no more attempts to mutate the profile
   disable_tracepoints(state);
 
-  reset_stats(state);
+  reset_stats_not_thread_safe(state);
 
   // Remove all state from the `Collectors::ThreadState` and connected downstream components
   rb_funcall(state->thread_context_collector_instance, rb_intern("reset_after_fork"), 0);
@@ -839,11 +929,15 @@ static VALUE _native_stats(DDTRACE_UNUSED VALUE self, VALUE instance) {
   struct cpu_and_wall_time_worker_state *state;
   TypedData_Get_Struct(instance, struct cpu_and_wall_time_worker_state, &cpu_and_wall_time_worker_typed_data, state);
 
-  VALUE pretty_sampling_time_ns_min = state->stats.sampling_time_ns_min == UINT64_MAX ? Qnil : ULL2NUM(state->stats.sampling_time_ns_min);
-  VALUE pretty_sampling_time_ns_max = state->stats.sampling_time_ns_max == 0 ? Qnil : ULL2NUM(state->stats.sampling_time_ns_max);
-  VALUE pretty_sampling_time_ns_total = state->stats.sampling_time_ns_total == 0 ? Qnil : ULL2NUM(state->stats.sampling_time_ns_total);
-  VALUE pretty_sampling_time_ns_avg =
-    state->stats.sampled == 0 ? Qnil : DBL2NUM(((double) state->stats.sampling_time_ns_total) / state->stats.sampled);
+  unsigned long total_cpu_samples_attempted = state->stats.cpu_sampled + state->stats.cpu_skipped;
+  VALUE effective_cpu_sample_rate =
+    total_cpu_samples_attempted == 0 ? Qnil : DBL2NUM(((double) state->stats.cpu_sampled) / total_cpu_samples_attempted);
+  unsigned long total_allocation_samples_attempted = state->stats.allocation_sampled + state->stats.allocation_skipped;
+  VALUE effective_allocation_sample_rate =
+    total_allocation_samples_attempted == 0 ? Qnil : DBL2NUM(((double) state->stats.allocation_sampled) / total_allocation_samples_attempted);
+
+  VALUE allocation_sampler_snapshot = state->allocation_profiling_enabled && state->dynamic_sampling_rate_enabled ?
+    discrete_dynamic_sampler_state_snapshot(&state->allocation_sampler) : Qnil;
 
   VALUE stats_as_hash = rb_hash_new();
   VALUE arguments[] = {
@@ -852,22 +946,43 @@ static VALUE _native_stats(DDTRACE_UNUSED VALUE self, VALUE instance) {
     ID2SYM(rb_intern("simulated_signal_delivery")),                  /* => */ UINT2NUM(state->stats.simulated_signal_delivery),
     ID2SYM(rb_intern("signal_handler_enqueued_sample")),             /* => */ UINT2NUM(state->stats.signal_handler_enqueued_sample),
     ID2SYM(rb_intern("signal_handler_wrong_thread")),                /* => */ UINT2NUM(state->stats.signal_handler_wrong_thread),
-    ID2SYM(rb_intern("sampled")),                                    /* => */ UINT2NUM(state->stats.sampled),
-    ID2SYM(rb_intern("skipped_sample_because_of_dynamic_sampling_rate")), /* => */ UINT2NUM(state->stats.skipped_sample_because_of_dynamic_sampling_rate),
     ID2SYM(rb_intern("postponed_job_skipped_already_existed")),      /* => */ UINT2NUM(state->stats.postponed_job_skipped_already_existed),
     ID2SYM(rb_intern("postponed_job_success")),                      /* => */ UINT2NUM(state->stats.postponed_job_success),
     ID2SYM(rb_intern("postponed_job_full")),                         /* => */ UINT2NUM(state->stats.postponed_job_full),
     ID2SYM(rb_intern("postponed_job_unknown_result")),               /* => */ UINT2NUM(state->stats.postponed_job_unknown_result),
-    ID2SYM(rb_intern("sampling_time_ns_min")),                       /* => */ pretty_sampling_time_ns_min,
-    ID2SYM(rb_intern("sampling_time_ns_max")),                       /* => */ pretty_sampling_time_ns_max,
-    ID2SYM(rb_intern("sampling_time_ns_total")),                     /* => */ pretty_sampling_time_ns_total,
-    ID2SYM(rb_intern("sampling_time_ns_avg")),                       /* => */ pretty_sampling_time_ns_avg,
-    ID2SYM(rb_intern("allocations_during_sample")),                  /* => */ UINT2NUM(state->stats.allocations_during_sample),
+    ID2SYM(rb_intern("interrupt_thread_attempts")),                  /* => */ UINT2NUM(state->stats.interrupt_thread_attempts),
+
+    // CPU Stats
+    ID2SYM(rb_intern("cpu_sampled")),                /* => */ UINT2NUM(state->stats.cpu_sampled),
+    ID2SYM(rb_intern("cpu_skipped")),                /* => */ UINT2NUM(state->stats.cpu_skipped),
+    ID2SYM(rb_intern("cpu_effective_sample_rate")),  /* => */ effective_cpu_sample_rate,
+    ID2SYM(rb_intern("cpu_sampling_time_ns_min")),   /* => */ RUBY_NUM_OR_NIL(state->stats.cpu_sampling_time_ns_min, != UINT64_MAX, ULL2NUM),
+    ID2SYM(rb_intern("cpu_sampling_time_ns_max")),   /* => */ RUBY_NUM_OR_NIL(state->stats.cpu_sampling_time_ns_max, > 0, ULL2NUM),
+    ID2SYM(rb_intern("cpu_sampling_time_ns_total")), /* => */ RUBY_NUM_OR_NIL(state->stats.cpu_sampling_time_ns_total, > 0, ULL2NUM),
+    ID2SYM(rb_intern("cpu_sampling_time_ns_avg")),   /* => */ RUBY_AVG_OR_NIL(state->stats.cpu_sampling_time_ns_total, state->stats.cpu_sampled),
+
+    // Allocation stats
+    ID2SYM(rb_intern("allocation_sampled")),                /* => */ state->allocation_profiling_enabled ? ULONG2NUM(state->stats.allocation_sampled) : Qnil,
+    ID2SYM(rb_intern("allocation_skipped")),                /* => */ state->allocation_profiling_enabled ? ULONG2NUM(state->stats.allocation_skipped) : Qnil,
+    ID2SYM(rb_intern("allocation_effective_sample_rate")),  /* => */ effective_allocation_sample_rate,
+    ID2SYM(rb_intern("allocation_sampling_time_ns_min")),   /* => */ RUBY_NUM_OR_NIL(state->stats.allocation_sampling_time_ns_min, != UINT64_MAX, ULL2NUM),
+    ID2SYM(rb_intern("allocation_sampling_time_ns_max")),   /* => */ RUBY_NUM_OR_NIL(state->stats.allocation_sampling_time_ns_max, > 0, ULL2NUM),
+    ID2SYM(rb_intern("allocation_sampling_time_ns_total")), /* => */ RUBY_NUM_OR_NIL(state->stats.allocation_sampling_time_ns_total, > 0, ULL2NUM),
+    ID2SYM(rb_intern("allocation_sampling_time_ns_avg")),   /* => */ RUBY_AVG_OR_NIL(state->stats.allocation_sampling_time_ns_total, state->stats.allocation_sampled),
+    ID2SYM(rb_intern("allocation_sampler_snapshot")),       /* => */ allocation_sampler_snapshot,
+    ID2SYM(rb_intern("allocations_during_sample")),         /* => */ state->allocation_profiling_enabled ? UINT2NUM(state->stats.allocations_during_sample) : Qnil,
   };
   for (long unsigned int i = 0; i < VALUE_COUNT(arguments); i += 2) rb_hash_aset(stats_as_hash, arguments[i], arguments[i+1]);
   return stats_as_hash;
 }
 
+static VALUE _native_stats_reset_not_thread_safe(DDTRACE_UNUSED VALUE self, VALUE instance) {
+  struct cpu_and_wall_time_worker_state *state;
+  TypedData_Get_Struct(instance, struct cpu_and_wall_time_worker_state, &cpu_and_wall_time_worker_typed_data, state);
+  reset_stats_not_thread_safe(state);
+  return Qnil;
+}
+
 void *simulate_sampling_signal_delivery(DDTRACE_UNUSED void *_unused) {
   struct cpu_and_wall_time_worker_state *state = active_sampler_instance_state; // Read from global variable, see "sampler global state safety" note above
 
@@ -885,9 +1000,17 @@ void *simulate_sampling_signal_delivery(DDTRACE_UNUSED void *_unused) {
 
 static void grab_gvl_and_sample(void) { rb_thread_call_with_gvl(simulate_sampling_signal_delivery, NULL); }
 
-static void reset_stats(struct cpu_and_wall_time_worker_state *state) {
-  state->stats = (struct stats) {}; // Resets all stats back to zero
-  state->stats.sampling_time_ns_min = UINT64_MAX; // Since we always take the min between existing and latest sample
+static void reset_stats_not_thread_safe(struct cpu_and_wall_time_worker_state *state) {
+  // NOTE: This is not really thread safe so ongoing sampling operations that are concurrent with a reset can have their stats:
+  //       * Lost (writes after stats retrieval but before reset).
+  //       * Included in the previous stats window (writes before stats retrieval and reset).
+  //       * Included in the following stats window (writes after stats retrieval and reset).
+  //       Given the expected infrequency of resetting (~once per 60s profile) and the auxiliary/non-critical nature of these stats
+  //       this momentary loss of accuracy is deemed acceptable to keep overhead to a minimum.
+  state->stats = (struct stats) {
+    .cpu_sampling_time_ns_min = UINT64_MAX, // Since we always take the min between existing and latest sample
+    .allocation_sampling_time_ns_min = UINT64_MAX, // Since we always take the min between existing and latest sample
+  };
 }
 
 static void sleep_for(uint64_t time_ns) {
@@ -937,8 +1060,16 @@ static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused)
     return;
   }
 
-  if (state->dynamic_sampling_rate_enabled && !discrete_dynamic_sampler_should_sample(&state->allocation_sampler)) {
-    return;
+  if (state->dynamic_sampling_rate_enabled) {
+    long now = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
+    if (now == 0) {
+      delayed_error(state, ERR_CLOCK_FAIL);
+      return;
+    }
+    if (!discrete_dynamic_sampler_should_sample(&state->allocation_sampler, now)) {
+      state->stats.allocation_skipped++;
+      return;
+    }
   }
 
   // @ivoanjo: Strictly speaking, this is not needed because Ruby should not call the same tracepoint while a previous
@@ -950,14 +1081,32 @@ static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused)
   // Rescue against any exceptions that happen during sampling
   safely_call(rescued_sample_allocation, tracepoint_data, state->self_instance);
 
-  discrete_dynamic_sampler_after_sample(&state->allocation_sampler);
+  if (state->dynamic_sampling_rate_enabled) {
+    long now = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
+    if (now == 0) {
+      delayed_error(state, ERR_CLOCK_FAIL);
+      // NOTE: Not short-circuiting here to make sure cleanup happens
+    }
+    uint64_t sampling_time_ns = discrete_dynamic_sampler_after_sample(&state->allocation_sampler, now);
+    // NOTE: To keep things lean when dynamic sampling rate is disabled we skip clock interactions which is
+    //       why we're fine with having this inside this conditional.
+    state->stats.allocation_sampling_time_ns_min = uint64_min_of(sampling_time_ns, state->stats.allocation_sampling_time_ns_min);
+    state->stats.allocation_sampling_time_ns_max = uint64_max_of(sampling_time_ns, state->stats.allocation_sampling_time_ns_max);
+    state->stats.allocation_sampling_time_ns_total += sampling_time_ns;
+  }
+
+  state->stats.allocation_sampled++;
 
   state->during_sample = false;
 }
 
 static void disable_tracepoints(struct cpu_and_wall_time_worker_state *state) {
-  rb_tracepoint_disable(state->gc_tracepoint);
-  rb_tracepoint_disable(state->object_allocation_tracepoint);
+  if (state->gc_tracepoint != Qnil) {
+    rb_tracepoint_disable(state->gc_tracepoint);
+  }
+  if (state->object_allocation_tracepoint != Qnil) {
+    rb_tracepoint_disable(state->object_allocation_tracepoint);
+  }
 }
 
 static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self) {
@@ -994,3 +1143,19 @@ static VALUE rescued_sample_allocation(VALUE tracepoint_data) {
   // Return a dummy VALUE because we're called from rb_rescue2 which requires it
   return Qnil;
 }
+
+static void delayed_error(struct cpu_and_wall_time_worker_state *state, const char *error) {
+  // If we can't raise an immediate exception at the calling site, use the asynchronous flow through the main worker loop.
+  stop_state(state, rb_exc_new_cstr(rb_eRuntimeError, error));
+}
+
+static VALUE _native_delayed_error(DDTRACE_UNUSED VALUE self, VALUE instance, VALUE error_msg) {
+  ENFORCE_TYPE(error_msg, T_STRING);
+
+  struct cpu_and_wall_time_worker_state *state;
+  TypedData_Get_Struct(instance, struct cpu_and_wall_time_worker_state, &cpu_and_wall_time_worker_typed_data, state);
+
+  delayed_error(state, rb_string_value_cstr(&error_msg));
+
+  return Qnil;
+}