ddtrace 1.12.1 → 1.23.2

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

@@ -12,10 +12,16 @@
 #include "collectors_thread_context.h"
 #include "collectors_dynamic_sampling_rate.h"
 #include "collectors_idle_sampling_helper.h"
+#include "collectors_discrete_dynamic_sampler.h"
 #include "private_vm_api_access.h"
 #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;
+
 // Used to trigger the execution of Collectors::ThreadState, which implements all of the sampling logic
 // itself; this class only implements the "when to do it" part.
 //
@@ -75,19 +81,28 @@
 //
 // ---
 
+#ifndef NO_POSTPONED_TRIGGER
+  // Used to call the rb_postponed_job_trigger from Ruby 3.3+. These get initialized in
+  // `collectors_cpu_and_wall_time_worker_init` below and always get reused after that.
+  static rb_postponed_job_handle_t sample_from_postponed_job_handle;
+  static rb_postponed_job_handle_t after_gc_from_postponed_job_handle;
+#endif
+
 // Contains state for a single CpuAndWallTimeWorker instance
 struct cpu_and_wall_time_worker_state {
   // These are immutable after initialization
 
   bool gc_profiling_enabled;
-  bool allocation_counting_enabled;
   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;
   VALUE owner_thread;
-  dynamic_sampling_rate_state dynamic_sampling_rate;
+  dynamic_sampling_rate_state cpu_dynamic_sampling_rate;
+  discrete_dynamic_sampler allocation_sampler;
   VALUE gc_tracepoint; // Used to get gc start/finish information
   VALUE object_allocation_tracepoint; // Used to get allocation counts and allocation profiling
 
@@ -107,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
@@ -117,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;
@@ -148,14 +177,17 @@ static VALUE _native_initialize(
   VALUE thread_context_collector_instance,
   VALUE gc_profiling_enabled,
   VALUE idle_sampling_helper_instance,
-  VALUE allocation_counting_enabled,
   VALUE no_signals_workaround_enabled,
-  VALUE dynamic_sampling_rate_enabled
+  VALUE dynamic_sampling_rate_enabled,
+  VALUE dynamic_sampling_rate_overhead_target_percentage,
+  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);
@@ -178,14 +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(DDTRACE_UNUSED VALUE tracepoint_data, DDTRACE_UNUSED void *unused);
+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:
 //
@@ -198,6 +234,11 @@ static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self);
 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.
 //
@@ -208,6 +249,18 @@ __thread uint64_t allocation_count = 0;
 void collectors_cpu_and_wall_time_worker_init(VALUE profiling_module) {
   rb_global_variable(&active_sampler_instance);
 
+  #ifndef NO_POSTPONED_TRIGGER
+    int unused_flags = 0;
+    sample_from_postponed_job_handle = rb_postponed_job_preregister(unused_flags, sample_from_postponed_job, NULL);
+    after_gc_from_postponed_job_handle = rb_postponed_job_preregister(unused_flags, after_gc_from_postponed_job, NULL);
+
+    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");
   VALUE collectors_cpu_and_wall_time_worker_class = rb_define_class_under(collectors_module, "CpuAndWallTimeWorker", rb_cObject);
   // Hosts methods used for testing the native code using RSpec
@@ -223,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, 7);
+  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);
@@ -239,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
@@ -257,14 +314,18 @@ static const rb_data_type_t cpu_and_wall_time_worker_typed_data = {
 static VALUE _native_new(VALUE klass) {
   struct cpu_and_wall_time_worker_state *state = ruby_xcalloc(1, sizeof(struct cpu_and_wall_time_worker_state));
 
+  // Note: Any exceptions raised from this note until the TypedData_Wrap_Struct call will lead to the state memory
+  // being leaked.
+
   state->gc_profiling_enabled = false;
-  state->allocation_counting_enabled = false;
   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->dynamic_sampling_rate);
+  dynamic_sampling_rate_init(&state->cpu_dynamic_sampling_rate);
   state->gc_tracepoint = Qnil;
   state->object_allocation_tracepoint = Qnil;
 
@@ -274,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);
 }
@@ -285,22 +354,39 @@ static VALUE _native_initialize(
   VALUE thread_context_collector_instance,
   VALUE gc_profiling_enabled,
   VALUE idle_sampling_helper_instance,
-  VALUE allocation_counting_enabled,
   VALUE no_signals_workaround_enabled,
-  VALUE dynamic_sampling_rate_enabled
+  VALUE dynamic_sampling_rate_enabled,
+  VALUE dynamic_sampling_rate_overhead_target_percentage,
+  VALUE allocation_profiling_enabled,
+  VALUE skip_idle_samples_for_testing
 ) {
   ENFORCE_BOOLEAN(gc_profiling_enabled);
-  ENFORCE_BOOLEAN(allocation_counting_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);
 
   state->gc_profiling_enabled = (gc_profiling_enabled == Qtrue);
-  state->allocation_counting_enabled = (allocation_counting_enabled == Qtrue);
   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) {
+    dynamic_sampling_rate_set_overhead_target_percentage(&state->cpu_dynamic_sampling_rate, total_overhead_target_percentage);
+  } else {
+    // 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);
+    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);
   state->idle_sampling_helper_instance = idle_sampling_helper_instance;
   state->gc_tracepoint = rb_tracepoint_new(Qnil, RUBY_INTERNAL_EVENT_GC_ENTER | RUBY_INTERNAL_EVENT_GC_EXIT, on_gc_event, NULL /* unused */);
@@ -327,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)) {
@@ -351,7 +443,9 @@ static VALUE _native_sampling_loop(DDTRACE_UNUSED VALUE _self, VALUE instance) {
   if (state->stop_thread == rb_thread_current()) return Qnil;
 
   // Reset the dynamic sampling rate state, if any (reminder: the monotonic clock reference may change after a fork)
-  dynamic_sampling_rate_reset(&state->dynamic_sampling_rate);
+  dynamic_sampling_rate_reset(&state->cpu_dynamic_sampling_rate);
+  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;
@@ -413,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;
 }
@@ -457,20 +555,50 @@ static void handle_sampling_signal(DDTRACE_UNUSED int _signal, DDTRACE_UNUSED si
 
   // Note: If we ever want to get rid of rb_postponed_job_register_one, remember not to clobber Ruby exceptions, as
   // this function does this helpful job for us now -- https://github.com/ruby/ruby/commit/a98e343d39c4d7bf1e2190b076720f32d9f298b3.
-  int result = rb_postponed_job_register_one(0, sample_from_postponed_job, 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
-  switch (result) {
-    case 0:
-      state->stats.postponed_job_full++; break;
-    case 1:
-      state->stats.postponed_job_success++; break;
-    case 2:
-      state->stats.postponed_job_skipped_already_existed++; break;
-    default:
-      state->stats.postponed_job_unknown_result++;
-  }
+  #ifndef NO_POSTPONED_TRIGGER // Ruby 3.3+
+    rb_postponed_job_trigger(sample_from_postponed_job_handle);
+    state->stats.postponed_job_success++; // Always succeeds
+  #else
+
+    // 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
+    switch (result) {
+      case 0:
+        state->stats.postponed_job_full++; break;
+      case 1:
+        state->stats.postponed_job_success++; break;
+      case 2:
+        state->stats.postponed_job_skipped_already_existed++; break;
+      default:
+        state->stats.postponed_job_unknown_result++;
+    }
+  #endif
 }
 
 // The actual sampling trigger loop always runs **without** the global vm lock.
@@ -498,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);
+        }
       }
     }
 
@@ -519,7 +653,7 @@ static void *run_sampling_trigger_loop(void *state_ptr) {
     // Note that we deliberately should NOT combine this sleep_for with the one above because the result of
     // `dynamic_sampling_rate_get_sleep` may have changed while the above sleep was ongoing.
     uint64_t extra_sleep =
-      dynamic_sampling_rate_get_sleep(&state->dynamic_sampling_rate, monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE));
+      dynamic_sampling_rate_get_sleep(&state->cpu_dynamic_sampling_rate, monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE));
     if (state->dynamic_sampling_rate_enabled && extra_sleep > 0) sleep_for(extra_sleep);
   }
 
@@ -559,12 +693,12 @@ 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->dynamic_sampling_rate, wall_time_ns_before_sample)) {
-    state->stats.skipped_sample_because_of_dynamic_sampling_rate++;
+  if (state->dynamic_sampling_rate_enabled && !dynamic_sampling_rate_should_sample(&state->cpu_dynamic_sampling_rate, wall_time_ns_before_sample)) {
+    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);
@@ -575,11 +709,11 @@ 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->dynamic_sampling_rate, wall_time_ns_after_sample, sampling_time_ns);
+  dynamic_sampling_rate_after_sample(&state->cpu_dynamic_sampling_rate, wall_time_ns_after_sample, sampling_time_ns);
 
   // Return a dummy VALUE because we're called from rb_rescue2 which requires it
   return Qnil;
@@ -617,7 +751,10 @@ static VALUE release_gvl_and_run_sampling_trigger_loop(VALUE instance) {
   // because they may raise exceptions.
   install_sigprof_signal_handler(handle_sampling_signal, "handle_sampling_signal");
   if (state->gc_profiling_enabled) rb_tracepoint_enable(state->gc_tracepoint);
-  if (state->allocation_counting_enabled) rb_tracepoint_enable(state->object_allocation_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);
 
@@ -699,28 +836,17 @@ static void on_gc_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused) {
   if (event == RUBY_INTERNAL_EVENT_GC_ENTER) {
     thread_context_collector_on_gc_start(state->thread_context_collector_instance);
   } else if (event == RUBY_INTERNAL_EVENT_GC_EXIT) {
-    // Design: In an earlier iteration of this feature (see https://github.com/DataDog/dd-trace-rb/pull/2308) we
-    // actually had a single method to implement the behavior of both thread_context_collector_on_gc_finish
-    // and thread_context_collector_sample_after_gc (the latter is called via after_gc_from_postponed_job).
-    //
-    // Unfortunately, then we discovered the safety issue around no allocations, and thus decided to separate them -- so that
-    // the sampling could run outside the tight safety constraints of the garbage collection process.
-    //
-    // There is a downside: The sample is now taken very very shortly afterwards the GC finishes, and not immediately
-    // as the GC finishes, which means the stack captured may by affected by "skid", e.g. point slightly after where
-    // it should be pointing at.
-    // Alternatives to solve this would be to capture no stack for garbage collection (as we do for Java and .net);
-    // making the sampling process allocation-safe (very hard); or separate stack sampling from sample recording,
-    // e.g. enabling us to capture the stack in thread_context_collector_on_gc_finish and do the rest later
-    // (medium hard).
-
-    thread_context_collector_on_gc_finish(state->thread_context_collector_instance);
-    // We use rb_postponed_job_register_one to ask Ruby to run thread_context_collector_sample_after_gc after if
-    // fully finishes the garbage collection, so that one is allowed to do allocations and throw exceptions as usual.
-    //
-    // Note: If we ever want to get rid of rb_postponed_job_register_one, remember not to clobber Ruby exceptions, as
-    // this function does this helpful job for us now -- https://github.com/ruby/ruby/commit/a98e343d39c4d7bf1e2190b076720f32d9f298b3.
-    rb_postponed_job_register_one(0, after_gc_from_postponed_job, NULL);
+    bool should_flush = thread_context_collector_on_gc_finish(state->thread_context_collector_instance);
+
+    // We use rb_postponed_job_register_one to ask Ruby to run thread_context_collector_sample_after_gc when the
+    // thread collector flags it's time to flush.
+    if (should_flush) {
+      #ifndef NO_POSTPONED_TRIGGER // Ruby 3.3+
+        rb_postponed_job_trigger(after_gc_from_postponed_job_handle);
+      #else
+        rb_postponed_job_register_one(0, after_gc_from_postponed_job, NULL);
+      #endif
+    }
   }
 }
 
@@ -787,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);
@@ -803,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[] = {
@@ -816,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
 
@@ -849,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) {
@@ -873,14 +1032,14 @@ static void sleep_for(uint64_t time_ns) {
 }
 
 static VALUE _native_allocation_count(DDTRACE_UNUSED VALUE self) {
-  bool is_profiler_running = active_sampler_instance_state != NULL;
+  bool are_allocations_being_tracked = active_sampler_instance_state != NULL && active_sampler_instance_state->allocation_profiling_enabled;
 
-  return is_profiler_running ? ULL2NUM(allocation_count) : Qnil;
+  return are_allocations_being_tracked ? ULL2NUM(allocation_count) : Qnil;
 }
 
 // Implements memory-related profiling events. This function is called by Ruby via the `object_allocation_tracepoint`
 // when the RUBY_INTERNAL_EVENT_NEWOBJ event is triggered.
-static void on_newobj_event(DDTRACE_UNUSED VALUE tracepoint_data, DDTRACE_UNUSED void *unused) {
+static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused) {
   // Update thread-local allocation count
   if (RB_UNLIKELY(allocation_count == UINT64_MAX)) {
     allocation_count = 0;
@@ -901,20 +1060,53 @@ static void on_newobj_event(DDTRACE_UNUSED VALUE tracepoint_data, DDTRACE_UNUSED
     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
   // invocation is still pending, (e.g. it wouldn't call `on_newobj_event` while it's already running), but I decided
   // to keep this here for consistency -- every call to the thread context (other than the special gc calls which are
   // defined as not being able to allocate) sets this.
   state->during_sample = true;
 
-  // TODO: Sampling goes here (calling into `thread_context_collector_sample_allocation`)
+  // Rescue against any exceptions that happen during sampling
+  safely_call(rescued_sample_allocation, tracepoint_data, state->self_instance);
+
+  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) {
@@ -929,3 +1121,41 @@ static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self) {
     return result;
   }
 }
+
+static VALUE rescued_sample_allocation(VALUE tracepoint_data) {
+  struct cpu_and_wall_time_worker_state *state = active_sampler_instance_state; // Read from global variable, see "sampler global state safety" note above
+
+  // This should not happen in a normal situation because on_newobj_event already checked for this, but just in case...
+  if (state == NULL) return Qnil;
+
+  rb_trace_arg_t *data = rb_tracearg_from_tracepoint(tracepoint_data);
+  VALUE new_object = rb_tracearg_object(data);
+
+  unsigned long allocations_since_last_sample = state->dynamic_sampling_rate_enabled ?
+    // if we're doing dynamic sampling, ask the sampler how many events since last sample
+    discrete_dynamic_sampler_events_since_last_sample(&state->allocation_sampler) :
+    // if we aren't, then we're sampling every event
+    1;
+  // TODO: Signal in the profile that clamping happened?
+  unsigned int weight = allocations_since_last_sample > MAX_ALLOC_WEIGHT ? MAX_ALLOC_WEIGHT : (unsigned int) allocations_since_last_sample;
+  thread_context_collector_sample_allocation(state->thread_context_collector_instance, weight, new_object);
+
+  // 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;
+}