datadog 2.36.0 → 2.37.0

data/ext/datadog_profiling_native_extension/collectors_thread_context.c

@@ -1,4 +1,5 @@
 #include <ruby.h>
+#include <ruby/debug.h>
 
 #include "datadog_ruby_common.h"
 #include "collectors_thread_context.h"
@@ -108,6 +109,15 @@ static ID server_id;      // id of :server in Ruby
 static ID otel_context_storage_id; // id of :__opentelemetry_context_storage__ in Ruby
 static ID otel_fiber_context_storage_id; // id of :@opentelemetry_context in Ruby
 
+// This is mutable and gets set last-writer-wins style whenever a new `ThreadContext` is created.
+//
+// The initial value should be kept in sync with the default for DD_PROFILING_MAX_FRAMES
+// in settings.rb. See `initialize_context` for details on why this is needed/used.
+static uint16_t latest_max_frames = 400;
+
+// Global tracepoint for RUBY_EVENT_THREAD_BEGIN. Created and enabled once when the first ThreadContext collector is initialized.
+static VALUE thread_begin_tracepoint = Qnil;
+
 
 typedef enum { OTEL_CONTEXT_ENABLED_FALSE, OTEL_CONTEXT_ENABLED_ONLY, OTEL_CONTEXT_ENABLED_BOTH } otel_context_enabled;
 typedef enum { OTEL_CONTEXT_SOURCE_UNKNOWN, OTEL_CONTEXT_SOURCE_FIBER_IVAR, OTEL_CONTEXT_SOURCE_FIBER_LOCAL } otel_context_source;
@@ -160,6 +170,7 @@ typedef struct {
     // How many per-thread samples were skipped because the thread has been continuously suspended
     // (no GVL) since its previous sample, so its Ruby stack cannot have changed.
     unsigned int inactive_thread_samples_skipped;
+    unsigned int profiler_thread_samples_skipped;
   } stats;
 
   struct {
@@ -171,7 +182,20 @@ typedef struct {
   } gc_tracking;
 } thread_context_collector_state;
 
-// Tracks per-thread state
+// Tracks per-thread state.
+// This state is global and lives forever on the Ruby Thread (until the Thread is GC'd).
+// The state is created early on for all threads on the main Ractor
+// (enabling a TracePoint only enables it for the current Ractor).
+// The state is either created when the Thread starts running (via the RUBY_EVENT_THREAD_BEGIN TracePoint),
+// or the first time we create a ThreadContext by iterating Thread.list.
+// Unfortunately that RUBY_EVENT_THREAD_BEGIN TracePoint still fires after some other events:
+// * RUBY_INTERNAL_THREAD_EVENT_RESUMED for the Thread acquiring the GVL for the first time
+// * an early SIGPROF calling handle_sampling_signal()
+// * Ruby might check for interrupts and run postponed jobs (e.g. thread_context_collector_sample)
+// * another RUBY_EVENT_THREAD_BEGIN TracePoint which might run before ours
+// * etc
+// For those cases we have to ignore those events and we cannot assume the state is always set,
+// however this only matters for a very short period when a thread starts.
 struct per_thread_context {
   sampling_buffer sampling_buffer;
   char thread_id[THREAD_ID_LIMIT_CHARS];
@@ -179,8 +203,8 @@ struct per_thread_context {
   char thread_invoke_location[THREAD_INVOKE_LOCATION_LIMIT_CHARS];
   ddog_CharSlice thread_invoke_location_char_slice;
   thread_cpu_time_id thread_cpu_time_id;
-  long cpu_time_at_previous_sample_ns;  // Can be INVALID_TIME until initialized or if getting it fails for another reason
-  long wall_time_at_previous_sample_ns; // Can be INVALID_TIME until initialized
+  long cpu_time_at_previous_sample_ns;
+  long wall_time_at_previous_sample_ns;
 
   // There are 3 possible states for the GVL (per thread), and 3 transitions for which we receive GVL events:
   // Thread holds the GVL
@@ -233,6 +257,14 @@ struct per_thread_context {
   // but this is deemed worth it for this optimization. In any case we don't know exactly
   // at what time a thread was doing CPU work (unless it's on CPU 100% of the time).
   bool was_skipped_at_last_sample;
+  // Set as true for CpuAndWallTimeWorker and IdleSamplingHelper threads.
+  // When true, per-tick samples are skipped entirely; the thread is sampled only once per
+  // reporting period during the on_serialize flush.
+  //
+  // These threads are always in native code so their stacks aren't interesting;
+  // the Profiling::Scheduler thread on the other hand does a lot of different
+  // things using a mix of Ruby and native code, so that one isn't considered internal.
+  bool is_profiler_internal_thread;
 
   struct {
     // Both of these fields are set by on_gc_start and kept until on_gc_finish is called.
@@ -270,7 +302,6 @@ static void update_metrics_and_sample(
   thread_context_collector_state *state,
   VALUE thread_being_sampled,
   per_thread_context *thread_context,
-  sampling_buffer* sampling_buffer,
   long current_cpu_time_ns,
   long current_monotonic_wall_time_ns,
   bool force_sample
@@ -279,7 +310,6 @@ static void trigger_sample_for_thread(
   thread_context_collector_state *state,
   VALUE thread_being_sampled,
   per_thread_context *thread_context,
-  sampling_buffer* sampling_buffer,
   sample_values values,
   long current_monotonic_wall_time_ns,
   ddog_CharSlice *ruby_vm_type,
@@ -288,8 +318,9 @@ static void trigger_sample_for_thread(
   bool is_safe_to_allocate_objects
 );
 static VALUE _native_thread_list(VALUE self);
-static per_thread_context *get_or_create_context_for(VALUE thread, thread_context_collector_state *state);
-static void initialize_context(VALUE thread, per_thread_context *thread_context, thread_context_collector_state *state);
+static void check_frozen_thread(VALUE thread);
+static per_thread_context *get_or_create_context_for(VALUE thread);
+static void initialize_context(VALUE thread, per_thread_context *thread_context);
 static VALUE _native_inspect(VALUE self, VALUE collector_instance);
 static VALUE per_thread_context_to_ruby_hash(per_thread_context *thread_context);
 static VALUE stats_to_ruby_hash(thread_context_collector_state *state, VALUE hash);
@@ -326,7 +357,6 @@ static bool handle_gvl_waiting(
   thread_context_collector_state *state,
   VALUE thread_being_sampled,
   per_thread_context *thread_context,
-  sampling_buffer* sampling_buffer,
   long current_cpu_time_ns
 );
 #ifndef NO_GVL_INSTRUMENTATION
@@ -335,8 +365,8 @@ static bool handle_gvl_waiting(
   static VALUE _native_on_gvl_running(DDTRACE_UNUSED VALUE self, VALUE collector_instance, VALUE thread);
   static VALUE _native_on_gvl_released(DDTRACE_UNUSED VALUE self, VALUE thread);
   static VALUE _native_sample_after_gvl_running(DDTRACE_UNUSED VALUE self, VALUE collector_instance, VALUE thread, VALUE allow_exception);
-  static VALUE _native_apply_delta_to_cpu_time_at_previous_sample_ns(DDTRACE_UNUSED VALUE self, VALUE thread, VALUE delta_ns);
 #endif
+static VALUE _native_apply_delta_to_cpu_time_at_previous_sample_ns(DDTRACE_UNUSED VALUE self, VALUE thread, VALUE delta_ns);
 static void otel_without_ddtrace_trace_identifiers_for(
   thread_context_collector_state *state,
   VALUE thread,
@@ -349,11 +379,15 @@ static VALUE safely_lookup_hash_without_going_into_ruby_code(VALUE hash, VALUE k
 static VALUE _native_system_epoch_time_now_ns(DDTRACE_UNUSED VALUE self, VALUE collector_instance);
 static VALUE _native_prepare_sample_inside_signal_handler(DDTRACE_UNUSED VALUE self);
 static VALUE _native_clear_per_thread_context_for(DDTRACE_UNUSED VALUE self, VALUE thread);
-static bool skip_sample(thread_context_collector_state *state, per_thread_context *thread_context, bool is_gvl_waiting_state, bool force_sample_suspended);
+static VALUE _native_mark_thread_as_profiler_internal(DDTRACE_UNUSED VALUE self, VALUE thread);
+static VALUE _native_remove_per_thread_context_for(DDTRACE_UNUSED VALUE self, VALUE thread);
+static bool skip_sample(thread_context_collector_state *state, per_thread_context *thread_context, bool is_gvl_waiting_state, bool force_sample);
+static void on_thread_begin_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused);
 
 void collectors_thread_context_init(VALUE profiling_module) {
   VALUE collectors_module = rb_define_module_under(profiling_module, "Collectors");
   VALUE collectors_thread_context_class = rb_define_class_under(collectors_module, "ThreadContext", rb_cObject);
+
   // Hosts methods used for testing the native code using RSpec
   VALUE testing_module = rb_define_module_under(collectors_thread_context_class, "Testing");
 
@@ -384,14 +418,16 @@ void collectors_thread_context_init(VALUE profiling_module) {
   rb_define_singleton_method(testing_module, "_native_system_epoch_time_now_ns", _native_system_epoch_time_now_ns, 1);
   rb_define_singleton_method(testing_module, "_native_prepare_sample_inside_signal_handler", _native_prepare_sample_inside_signal_handler, 0);
   rb_define_singleton_method(testing_module, "_native_clear_per_thread_context_for", _native_clear_per_thread_context_for, 1);
+  rb_define_singleton_method(testing_module, "_native_remove_per_thread_context_for", _native_remove_per_thread_context_for, 1);
+  rb_define_singleton_method(testing_module, "_native_mark_thread_as_profiler_internal", _native_mark_thread_as_profiler_internal, 1);
   #ifndef NO_GVL_INSTRUMENTATION
     rb_define_singleton_method(testing_module, "_native_on_gvl_waiting", _native_on_gvl_waiting, 1);
     rb_define_singleton_method(testing_module, "_native_gvl_waiting_at_for", _native_gvl_waiting_at_for, 1);
     rb_define_singleton_method(testing_module, "_native_on_gvl_running", _native_on_gvl_running, 2);
     rb_define_singleton_method(testing_module, "_native_on_gvl_released", _native_on_gvl_released, 1);
     rb_define_singleton_method(testing_module, "_native_sample_after_gvl_running", _native_sample_after_gvl_running, 3);
-    rb_define_singleton_method(testing_module, "_native_apply_delta_to_cpu_time_at_previous_sample_ns", _native_apply_delta_to_cpu_time_at_previous_sample_ns, 2);
   #endif
+  rb_define_singleton_method(testing_module, "_native_apply_delta_to_cpu_time_at_previous_sample_ns", _native_apply_delta_to_cpu_time_at_previous_sample_ns, 2);
 
   at_active_span_id = rb_intern_const("@active_span");
   at_active_trace_id = rb_intern_const("@active_trace");
@@ -417,6 +453,8 @@ void collectors_thread_context_init(VALUE profiling_module) {
   // This will raise if Ruby already ran out of thread-local keys
   per_thread_context_tls_init();
 
+  rb_global_variable(&thread_begin_tracepoint);
+
   gc_profiling_init();
 }
 
@@ -485,13 +523,20 @@ static void per_thread_context_typed_data_free(void *ctx_ptr) {
   free(ctx);
 }
 
-static VALUE _native_clear_per_thread_context_for(DDTRACE_UNUSED VALUE self, VALUE thread) {
+static VALUE _native_clear_per_thread_context_for(VALUE self, VALUE thread) {
+  _native_remove_per_thread_context_for(self, thread);
+  get_or_create_context_for(thread);
+  return Qnil;
+}
+
+// Only for testing: removes the per-thread context without recreating it, so the thread has no context.
+// This simulates a thread that starts running before the RUBY_EVENT_THREAD_BEGIN tracepoint fires.
+static VALUE _native_remove_per_thread_context_for(DDTRACE_UNUSED VALUE self, VALUE thread) {
+  check_frozen_thread(thread);
   per_thread_context *ctx = get_per_thread_context(thread);
   if (ctx != NULL) {
     set_per_thread_context(thread, NULL);
-    if (!RB_OBJ_FROZEN(thread)) {
-      rb_ivar_set(thread, dd_per_thread_context_id, Qnil);
-    }
+    rb_ivar_set(thread, dd_per_thread_context_id, Qnil);
   }
   return Qnil;
 }
@@ -554,12 +599,15 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   ENFORCE_BOOLEAN(native_filenames_enabled);
   ENFORCE_TYPE(overhead_filename, T_STRING);
 
+  uint16_t max_frame_int = sampling_buffer_check_max_frames(NUM2INT(max_frames));
+  latest_max_frames = max_frame_int;
+
   thread_context_collector_state *state;
   TypedData_Get_Struct(self_instance, thread_context_collector_state, &thread_context_collector_typed_data, state);
 
   // Update this when modifying state struct
-  state->locations.len = sampling_buffer_check_max_frames(NUM2INT(max_frames));
-  state->locations.ptr = ruby_xcalloc(state->locations.len, sizeof(ddog_prof_Location));
+  state->locations.len = max_frame_int;
+  state->locations.ptr = ruby_xcalloc(max_frame_int, sizeof(ddog_prof_Location));
   state->recorder_instance = enforce_recorder_instance(recorder_instance);
   recorder_install_on_serialize(recorder_instance, self_instance);
   state->endpoint_collection_enabled = (endpoint_collection_enabled == Qtrue);
@@ -585,6 +633,19 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
     state->tracer_context_key = rb_to_id(tracer_context_key);
   }
 
+  if (thread_begin_tracepoint == Qnil) {
+    thread_begin_tracepoint = rb_tracepoint_new(Qnil, RUBY_EVENT_THREAD_BEGIN, on_thread_begin_event, NULL);
+    rb_tracepoint_enable(thread_begin_tracepoint);
+
+    VALUE thread_list = rb_ary_new();
+    ddtrace_thread_list(thread_list);
+    long thread_count = RARRAY_LEN(thread_list);
+    for (long i = 0; i < thread_count; i++) {
+      get_or_create_context_for(RARRAY_AREF(thread_list, i));
+    }
+    RB_GC_GUARD(thread_list);
+  }
+
   return Qtrue;
 }
 
@@ -698,7 +759,7 @@ void thread_context_collector_sample(VALUE self_instance, long current_monotonic
   TypedData_Get_Struct(self_instance, thread_context_collector_state, &thread_context_collector_typed_data, state);
 
   VALUE current_thread = rb_thread_current();
-  per_thread_context *current_thread_context = get_or_create_context_for(current_thread, state);
+  per_thread_context *current_thread_context = get_or_create_context_for(current_thread);
   long cpu_time_at_sample_start_for_current_thread = cpu_time_now_ns(current_thread_context);
 
   VALUE threads = thread_list(state);
@@ -706,7 +767,7 @@ void thread_context_collector_sample(VALUE self_instance, long current_monotonic
   const long thread_count = RARRAY_LEN(threads);
   for (long i = 0; i < thread_count; i++) {
     VALUE thread = RARRAY_AREF(threads, i);
-    per_thread_context *thread_context = get_or_create_context_for(thread, state);
+    per_thread_context *thread_context = get_or_create_context_for(thread);
 
     // We account for cpu-time for the current thread in a different way: we use the cpu-time at sampling start,
     // to avoid blaming the time the profiler took on whatever is currently running on the thread,
@@ -717,29 +778,36 @@ void thread_context_collector_sample(VALUE self_instance, long current_monotonic
       state,
       thread,
       thread_context,
-      &thread_context->sampling_buffer,
       current_cpu_time_ns,
       current_monotonic_wall_time_ns,
       false);
   }
 
   state->stats.sample_count++;
-  record_sampling_overhead(state, current_thread_context);
+
+  // If the current thread is a profiler-internal thread, we don't use record_sampling_overhead()
+  // and accept the sampling overhead is attributed to the profiler-internal thread
+  // (both are under the same datadog group in the flamegraph).
+  // The reason we need to do this is profiler-internal threads are skipped during per-tick sampling
+  // so their timestamps are not updated. record_sampling_overhead() would see the stale previous-sample timestamps and
+  // attribute the entire sleep interval incorrectly as overhead instead of just the sampling work.
+  if (!current_thread_context->is_profiler_internal_thread) {
+    record_sampling_overhead(state, current_thread_context);
+  }
 }
 
 static void update_metrics_and_sample(
   thread_context_collector_state *state,
   VALUE thread_being_sampled,
   per_thread_context *thread_context,
-  sampling_buffer* sampling_buffer,
   long current_cpu_time_ns,
   long current_monotonic_wall_time_ns,
-  bool force_sample_suspended
+  bool force_sample
 ) {
   bool is_gvl_waiting_state =
-    handle_gvl_waiting(state, thread_being_sampled, thread_context, sampling_buffer, current_cpu_time_ns);
+    handle_gvl_waiting(state, thread_being_sampled, thread_context, current_cpu_time_ns);
 
-  if (skip_sample(state, thread_context, is_gvl_waiting_state, force_sample_suspended)) return;
+  if (skip_sample(state, thread_context, is_gvl_waiting_state, force_sample)) return;
 
   // Don't assign/update cpu during "Waiting for GVL"
   long cpu_time_elapsed_ns = is_gvl_waiting_state ? 0 : update_time_since_previous_sample(
@@ -775,7 +843,6 @@ static void update_metrics_and_sample(
     state,
     thread_being_sampled,
     thread_context,
-    sampling_buffer,
     (sample_values) {.cpu_time_ns = cpu_time_elapsed_ns, .cpu_or_wall_samples = 1, .wall_time_ns = wall_time_elapsed_ns},
     current_monotonic_wall_time_ns,
     NULL,
@@ -785,7 +852,12 @@ static void update_metrics_and_sample(
   );
 }
 
-static bool skip_sample(thread_context_collector_state *state, per_thread_context *thread_context, bool is_gvl_waiting_state, bool force_sample_suspended) {
+static bool skip_sample(thread_context_collector_state *state, per_thread_context *thread_context, bool is_gvl_waiting_state, bool force_sample) {
+  if (!force_sample && thread_context->is_profiler_internal_thread) {
+    state->stats.profiler_thread_samples_skipped++;
+    return true;
+  }
+
   // Racy read but harmless, can only cause an extra sample
   uint64_t gvl_state_change_count = thread_context->gvl_state_change_count;
 
@@ -797,7 +869,7 @@ static bool skip_sample(thread_context_collector_state *state, per_thread_contex
   // in handle_gvl_waiting (situation 1 extra sample, situation 2 regular sample) keeps running.
   // TODO: we could probably also skip while "Waiting for GVL"
   if (!is_gvl_waiting_state &&
-      !force_sample_suspended &&
+      !force_sample &&
       (gvl_state_change_count & GVL_SUSPENDED) &&
       gvl_state_change_count == thread_context->gvl_state_change_count_at_previous_sample) {
     state->stats.inactive_thread_samples_skipped++;
@@ -829,10 +901,8 @@ void thread_context_collector_on_gc_start(VALUE self_instance) {
 
   per_thread_context *thread_context = get_per_thread_context(rb_thread_current());
 
-  // If there was no previously-existing context for this thread, we won't allocate one (see safety). For now we just drop
-  // the GC sample, under the assumption that "a thread that is so new that we never sampled it even once before it triggers
-  // GC" is a rare enough case that we can just ignore it.
-  // We can always improve this later if we find that this happens often (and we have the counter to help us figure that out)!
+  // Context is created eagerly via on_thread_begin_event, so this should not normally be NULL.
+  // We can't get_or_create_context_for() here since we can't allocate (GC context).
   if (thread_context == NULL) {
     state->stats.gc_samples_missed_due_to_missing_context++;
     return;
@@ -862,8 +932,7 @@ bool thread_context_collector_on_gc_finish(VALUE self_instance) {
 
   per_thread_context *thread_context = get_per_thread_context(rb_thread_current());
 
-  // If there was no previously-existing context for this thread, we won't allocate one (see safety). We keep a metric for
-  // how often this happens -- see on_gc_start.
+  // Context is created eagerly, so this should not normally be NULL (see on_gc_start).
   if (thread_context == NULL) return false;
 
   long cpu_time_at_start_ns = thread_context->gc_tracking.cpu_time_at_start_ns;
@@ -982,7 +1051,6 @@ static void trigger_sample_for_thread(
   thread_context_collector_state *state,
   VALUE thread_being_sampled,
   per_thread_context *thread_context,
-  sampling_buffer* sampling_buffer,
   sample_values values,
   long current_monotonic_wall_time_ns,
   // These two labels are only used for allocation profiling; @ivoanjo: may want to refactor this at some point?
@@ -1104,7 +1172,7 @@ static void trigger_sample_for_thread(
 
   sample_thread(
     thread_being_sampled,
-    sampling_buffer,
+    &thread_context->sampling_buffer,
     state->locations,
     state->recorder_instance,
     values,
@@ -1133,17 +1201,22 @@ static VALUE _native_thread_list(DDTRACE_UNUSED VALUE _self) {
   return result;
 }
 
+static void check_frozen_thread(VALUE thread) {
+  if (RB_OBJ_FROZEN(thread)) {
+    raise_error(rb_eFrozenError, "Cannot setup profiler state for Thread %"PRIsVALUE" because it is frozen. Please avoid freezing Thread instances and/or report the issue to dd-trace-rb", thread);
+  }
+}
+
+// See the docs on struct per_thread_context.
 // This allocates a Ruby object and therefore needs the GVL and is not safe to call from RUBY_INTERNAL_EVENT_* hooks.
-static per_thread_context *get_or_create_context_for(VALUE thread, thread_context_collector_state *state) {
+static per_thread_context *get_or_create_context_for(VALUE thread) {
   per_thread_context *thread_context = get_per_thread_context(thread);
   if (thread_context != NULL) return thread_context;
 
-  if (RB_OBJ_FROZEN(thread)) {
-    raise_error(rb_eFrozenError, "Cannot setup profiler state for Thread %"PRIsVALUE" because it is frozen. Please avoid freezing Thread instances and/or report the issue to dd-trace-rb", thread);
-  }
+  check_frozen_thread(thread);
 
   thread_context = calloc(1, sizeof(per_thread_context)); // See "note on calloc vs ruby_xcalloc use" in heap_recorder.c
-  initialize_context(thread, thread_context, state);
+  initialize_context(thread, thread_context);
 
   VALUE wrapper = TypedData_Wrap_Struct(rb_cObject, &per_thread_context_typed_data, thread_context);
   rb_ivar_set(thread, dd_per_thread_context_id, wrapper);
@@ -1152,6 +1225,14 @@ static per_thread_context *get_or_create_context_for(VALUE thread, thread_contex
   return thread_context;
 }
 
+static void on_thread_begin_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused) {
+  if (!ddtrace_rb_ractor_main_p()) return;
+
+  VALUE thread = rb_tracearg_self(rb_tracearg_from_tracepoint(tracepoint_data));
+  ENFORCE_THREAD(thread);
+  get_or_create_context_for(thread);
+}
+
 #define LOGGING_GEM_PATH "/lib/logging/diagnostic_context.rb"
 
 // The `logging` gem monkey patches thread creation, which makes the `invoke_location_for` useless, since every thread
@@ -1173,8 +1254,12 @@ static bool is_logging_gem_monkey_patch(VALUE invoke_file_location) {
   return strncmp(invoke_file + invoke_file_len - logging_gem_path_len, LOGGING_GEM_PATH, logging_gem_path_len) == 0;
 }
 
-static void initialize_context(VALUE thread, per_thread_context *thread_context, thread_context_collector_state *state) {
-  sampling_buffer_initialize(&thread_context->sampling_buffer, state->locations.len);
+static void initialize_context(VALUE thread, per_thread_context *thread_context) {
+  // We always create per_thread_context's with latest_max_frames because
+  // 1) we don't always have access to the ThreadContext object (e.g. in a TracePoint).
+  // 2) the per_thread_context's are global and so might not match the ThreadContext#max_frames anyway.
+  // This is fine because sample_thread() handles when they don't match and resizes as needed.
+  sampling_buffer_initialize(&thread_context->sampling_buffer, latest_max_frames);
 
   snprintf(thread_context->thread_id, THREAD_ID_LIMIT_CHARS, "%"PRIu64" (%lu)", native_thread_id_for(thread), (unsigned long) thread_id_for(thread));
   thread_context->thread_id_char_slice = (ddog_CharSlice) {.ptr = thread_context->thread_id, .len = strlen(thread_context->thread_id)};
@@ -1193,7 +1278,7 @@ static void initialize_context(VALUE thread, per_thread_context *thread_context,
     } else {
       snprintf(thread_context->thread_invoke_location, THREAD_INVOKE_LOCATION_LIMIT_CHARS, "%s", "(Unnamed thread)");
     }
-  } else if (thread != state->main_thread) {
+  } else if (thread != rb_thread_main()) {
     // If the first function of a thread is native code, there won't be an invoke location, so we use this fallback.
     // NOTE: In the future, I wonder if we could take the pointer to the native function, and try to see if there's a native
     // symbol attached to it.
@@ -1207,9 +1292,8 @@ static void initialize_context(VALUE thread, per_thread_context *thread_context,
 
   thread_context->thread_cpu_time_id = thread_cpu_time_id_for(thread);
 
-  // These will get initialized during actual sampling
-  thread_context->cpu_time_at_previous_sample_ns = INVALID_TIME;
-  thread_context->wall_time_at_previous_sample_ns = INVALID_TIME;
+  thread_context->wall_time_at_previous_sample_ns = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
+  thread_context->cpu_time_at_previous_sample_ns = cpu_time_now_ns(thread_context);
 
   // These will only be used during a GC operation
   thread_context->gc_tracking.cpu_time_at_start_ns = INVALID_TIME;
@@ -1271,6 +1355,7 @@ static VALUE per_thread_context_to_ruby_hash(per_thread_context *thread_context)
     ID2SYM(rb_intern("gvl_state_change_count")), /* => */ ULL2NUM(thread_context->gvl_state_change_count),
     ID2SYM(rb_intern("gvl_state_change_count_at_previous_sample")), /* => */ ULL2NUM(thread_context->gvl_state_change_count_at_previous_sample),
     ID2SYM(rb_intern("was_skipped_at_last_sample")), /* => */ thread_context->was_skipped_at_last_sample ? Qtrue : Qfalse,
+    ID2SYM(rb_intern("is_profiler_internal_thread")), /* => */ thread_context->is_profiler_internal_thread ? Qtrue : Qfalse,
   };
   for (long unsigned int i = 0; i < VALUE_COUNT(arguments); i += 2) rb_hash_aset(context_as_hash, arguments[i], arguments[i+1]);
 
@@ -1284,6 +1369,7 @@ static VALUE stats_to_ruby_hash(thread_context_collector_state *state, VALUE has
     ID2SYM(rb_intern("gc_samples")),                               /* => */ UINT2NUM(state->stats.gc_samples),
     ID2SYM(rb_intern("gc_samples_missed_due_to_missing_context")), /* => */ UINT2NUM(state->stats.gc_samples_missed_due_to_missing_context),
     ID2SYM(rb_intern("inactive_thread_samples_skipped")),          /* => */ UINT2NUM(state->stats.inactive_thread_samples_skipped),
+    ID2SYM(rb_intern("profiler_thread_samples_skipped")),          /* => */ UINT2NUM(state->stats.profiler_thread_samples_skipped),
   };
   for (long unsigned int i = 0; i < VALUE_COUNT(arguments); i += 2) rb_hash_aset(hash, arguments[i], arguments[i+1]);
   return hash;
@@ -1634,11 +1720,10 @@ bool thread_context_collector_sample_allocation(VALUE self_instance, per_thread_
     state,
     current_thread,
     thread_context,
-    &thread_context->sampling_buffer,
     (sample_values) {.alloc_samples = sample_weight, .alloc_samples_unscaled = 1, .heap_sample = true},
     INVALID_TIME, // For now we're not collecting timestamps for allocation events, as per profiling team internal discussions
     &ruby_vm_type,
-     &class_name,
+    &class_name,
     /* is_gvl_waiting_state: */ false,
     /* is_safe_to_allocate_objects: */ false // Not safe to allocate further inside the NEWOBJ tracepoint
   );
@@ -1649,9 +1734,10 @@ bool thread_context_collector_sample_allocation(VALUE self_instance, per_thread_
 // This method exists only to enable testing Datadog::Profiling::Collectors::ThreadContext behavior using RSpec.
 // It SHOULD NOT be used for other purposes.
 static VALUE _native_sample_allocation(DDTRACE_UNUSED VALUE self, VALUE collector_instance, VALUE sample_weight, VALUE new_object) {
-  thread_context_collector_state *state;
-  TypedData_Get_Struct(collector_instance, thread_context_collector_state, &thread_context_collector_typed_data, state);
-  per_thread_context *thread_context = get_or_create_context_for(rb_thread_current(), state);
+  per_thread_context *thread_context = get_per_thread_context(rb_thread_current());
+  if (thread_context == NULL) {
+    rb_raise(rb_eRuntimeError, "Missing per_thread_context for current thread in _native_sample_allocation");
+  }
 
   debug_enter_unsafe_context();
 
@@ -1972,41 +2058,81 @@ void thread_context_collector_stats_reset_not_thread_safe(VALUE self_instance) {
   state->stats = (struct stats) {};
 }
 
-#ifndef NO_GVL_INSTRUMENTATION
-  void thread_context_collector_on_gvl_released(per_thread_context *thread_context) {
-    thread_context->gvl_state_change_count |= GVL_SUSPENDED;
+static void mark_thread_as_profiler_internal(per_thread_context *ctx) {
+  ctx->is_profiler_internal_thread = true;
+}
+
+void thread_context_collector_profiler_internal_thread_started(void) {
+  per_thread_context *ctx = get_or_create_context_for(rb_thread_current());
+  mark_thread_as_profiler_internal(ctx);
+}
+
+static VALUE _native_mark_thread_as_profiler_internal(DDTRACE_UNUSED VALUE self, VALUE thread) {
+  per_thread_context *ctx = get_or_create_context_for(thread);
+  mark_thread_as_profiler_internal(ctx);
+  return Qnil;
+}
+
+// Called via rb_ensure when a profiler-internal thread (worker or idle helper) is about to exit.
+// Records a final sample so the thread's accumulated cpu/wall time since the last on_serialize
+// flush is not lost. on_serialize (below) also flushes profiler-internal threads during periodic
+// serialization, but it can't help at shutdown: by the time the final serialize runs, these
+// threads are already dead and absent from thread_list.
+void thread_context_collector_profiler_internal_thread_done(VALUE self_instance) {
+  thread_context_collector_state *state;
+  TypedData_Get_Struct(self_instance, thread_context_collector_state, &thread_context_collector_typed_data, state);
+
+  VALUE current_thread = rb_thread_current();
+  per_thread_context *thread_context = get_or_create_context_for(current_thread);
+  if (!thread_context->is_profiler_internal_thread) {
+    rb_raise(rb_eRuntimeError, "current thread %"PRIsVALUE" is not profiler-internal thread", current_thread);
   }
 
-  // Called by the stack recorder at the start of _native_serialize, so that threads whose last
-  // per-tick sample was skipped by the SUSPENDED-skip optimization still get their accumulated
-  // time recorded in this reporting period. Without this, a thread that sleeps across the whole
-  // period would not be reported at all.
-  void thread_context_collector_on_serialize(VALUE self_instance) {
-    thread_context_collector_state *state;
-    TypedData_Get_Struct(self_instance, thread_context_collector_state, &thread_context_collector_typed_data, state);
+  long current_cpu_time_ns = cpu_time_now_ns(thread_context);
+  long current_monotonic_wall_time_ns = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
 
-    long current_monotonic_wall_time_ns = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
-    VALUE threads = thread_list(state);
-    const long thread_count = RARRAY_LEN(threads);
+  update_metrics_and_sample(
+    state,
+    current_thread,
+    thread_context,
+    current_cpu_time_ns,
+    current_monotonic_wall_time_ns,
+    true);
+}
 
-    for (long i = 0; i < thread_count; i++) {
-      VALUE thread = RARRAY_AREF(threads, i);
-      per_thread_context *thread_context = get_per_thread_context(thread);
-
-      if (thread_context != NULL && thread_context->was_skipped_at_last_sample) {
-        long current_cpu_time_ns = cpu_time_now_ns(thread_context);
-        // We need to force_sample_suspended=true otherwise this sample would be skipped too
-        update_metrics_and_sample(
-          state,
-          thread,
-          thread_context,
-          &thread_context->sampling_buffer,
-          current_cpu_time_ns,
-          current_monotonic_wall_time_ns,
-          true);
-      }
+// Flushes threads whose last per-tick sample was skipped (either by the SUSPENDED-skip
+// optimization, or by is_profiler_internal_thread) so their accumulated time is recorded.
+// Called by the stack recorder at the start of _native_serialize (regular periodic flush).
+void thread_context_collector_on_serialize(VALUE self_instance) {
+  thread_context_collector_state *state;
+  TypedData_Get_Struct(self_instance, thread_context_collector_state, &thread_context_collector_typed_data, state);
+
+  long current_monotonic_wall_time_ns = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
+  VALUE threads = thread_list(state);
+  const long thread_count = RARRAY_LEN(threads);
+
+  for (long i = 0; i < thread_count; i++) {
+    VALUE thread = RARRAY_AREF(threads, i);
+    per_thread_context *thread_context = get_per_thread_context(thread);
+
+    if (thread_context != NULL && (thread_context->was_skipped_at_last_sample || thread_context->is_profiler_internal_thread)) {
+      long current_cpu_time_ns = cpu_time_now_ns(thread_context);
+      // We need to force_sample=true otherwise this sample would be skipped too
+      update_metrics_and_sample(
+        state,
+        thread,
+        thread_context,
+        current_cpu_time_ns,
+        current_monotonic_wall_time_ns,
+        true);
     }
   }
+}
+
+#ifndef NO_GVL_INSTRUMENTATION
+  void thread_context_collector_on_gvl_released(per_thread_context *thread_context) {
+    thread_context->gvl_state_change_count |= GVL_SUSPENDED;
+  }
 
   void thread_context_collector_on_gvl_waiting(per_thread_context *thread_context) {
     long current_monotonic_wall_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
@@ -2103,7 +2229,7 @@ void thread_context_collector_stats_reset_not_thread_safe(VALUE self_instance) {
     thread_context_collector_state *state;
     TypedData_Get_Struct(self_instance, thread_context_collector_state, &thread_context_collector_typed_data, state);
 
-    per_thread_context *thread_context = get_or_create_context_for(current_thread, state);
+    per_thread_context *thread_context = get_or_create_context_for(current_thread);
 
     long gvl_waiting_at = thread_context->gvl_waiting_at;
 
@@ -2126,7 +2252,6 @@ void thread_context_collector_stats_reset_not_thread_safe(VALUE self_instance) {
       state,
       current_thread,
       thread_context,
-      &thread_context->sampling_buffer,
       cpu_time_for_thread,
       current_monotonic_wall_time_ns,
       false);
@@ -2141,7 +2266,6 @@ void thread_context_collector_stats_reset_not_thread_safe(VALUE self_instance) {
     thread_context_collector_state *state,
     VALUE thread_being_sampled,
     per_thread_context *thread_context,
-    sampling_buffer* sampling_buffer,
     long current_cpu_time_ns
   ) {
     long gvl_waiting_at = thread_context->gvl_waiting_at;
@@ -2219,7 +2343,6 @@ void thread_context_collector_stats_reset_not_thread_safe(VALUE self_instance) {
         state,
         thread_being_sampled,
         thread_context,
-        sampling_buffer,
         (sample_values) {.cpu_time_ns = cpu_time_elapsed_ns, .cpu_or_wall_samples = 1, .wall_time_ns = duration_until_start_of_gvl_waiting_ns},
         gvl_waiting_started_wall_time_ns,
         NULL,
@@ -2306,29 +2429,27 @@ void thread_context_collector_stats_reset_not_thread_safe(VALUE self_instance) {
     return result;
   }
 
-  static VALUE _native_apply_delta_to_cpu_time_at_previous_sample_ns(DDTRACE_UNUSED VALUE self, VALUE thread, VALUE delta_ns) {
-    ENFORCE_THREAD(thread);
-
-    per_thread_context *thread_context = get_per_thread_context(thread);
-    if (thread_context == NULL) raise_error(rb_eArgError, "Unexpected: This method cannot be used unless the per-thread context for the thread already exists");
-
-    thread_context->cpu_time_at_previous_sample_ns += NUM2LONG(delta_ns);
-
-    return Qtrue;
-  }
-
 #else
   static bool handle_gvl_waiting(
     DDTRACE_UNUSED thread_context_collector_state *state,
     DDTRACE_UNUSED VALUE thread_being_sampled,
     DDTRACE_UNUSED per_thread_context *thread_context,
-    DDTRACE_UNUSED sampling_buffer* sampling_buffer,
     DDTRACE_UNUSED long current_cpu_time_ns
   ) { return false; }
 
-  void thread_context_collector_on_serialize(DDTRACE_UNUSED VALUE self_instance) { }
 #endif // NO_GVL_INSTRUMENTATION
 
+static VALUE _native_apply_delta_to_cpu_time_at_previous_sample_ns(DDTRACE_UNUSED VALUE self, VALUE thread, VALUE delta_ns) {
+  ENFORCE_THREAD(thread);
+
+  per_thread_context *thread_context = get_per_thread_context(thread);
+  if (thread_context == NULL) raise_error(rb_eArgError, "Unexpected: This method cannot be used unless the per-thread context for the thread already exists");
+
+  thread_context->cpu_time_at_previous_sample_ns += NUM2LONG(delta_ns);
+
+  return Qtrue;
+}
+
 #define MAX_SAFE_LOOKUP_SIZE 16
 
 typedef struct { VALUE lookup_key; VALUE result; } safe_lookup_hash_state;