datadog 2.4.0 → 2.44.0

data/ext/datadog_profiling_native_extension/stack_recorder.c

@@ -8,6 +8,8 @@
 #include "ruby_helpers.h"
 #include "time_helpers.h"
 #include "heap_recorder.h"
+#include "encoded_profile.h"
+#include "collectors_thread_context.h"
 
 // Used to wrap a ddog_prof_Profile in a Ruby object and expose Ruby-level serialization APIs
 // This file implements the native bits of the Datadog::Profiling::StackRecorder class
@@ -130,69 +132,65 @@
 static VALUE ok_symbol = Qnil; // :ok in Ruby
 static VALUE error_symbol = Qnil; // :error in Ruby
 
-// Note: Please DO NOT use `VALUE_STRING` anywhere else, instead use `DDOG_CHARSLICE_C`.
-// `VALUE_STRING` is only needed because older versions of gcc (4.9.2, used in our Ruby 2.2 CI test images)
-// tripped when compiling `enabled_value_types` using `-std=gnu99` due to the extra cast that is included in
-// `DDOG_CHARSLICE_C` with the following error:
-//
-// ```
-// compiling ../../../../ext/ddtrace_profiling_native_extension/stack_recorder.c
-// ../../../../ext/ddtrace_profiling_native_extension/stack_recorder.c:23:1: error: initializer element is not constant
-// static const ddog_prof_ValueType enabled_value_types[] = {CPU_TIME_VALUE, CPU_SAMPLES_VALUE, WALL_TIME_VALUE};
-// ^
-// ```
-#define VALUE_STRING(string) {.ptr = "" string, .len = sizeof(string) - 1}
-
-#define CPU_TIME_VALUE          {.type_ = VALUE_STRING("cpu-time"),          .unit = VALUE_STRING("nanoseconds")}
 #define CPU_TIME_VALUE_ID 0
-#define CPU_SAMPLES_VALUE       {.type_ = VALUE_STRING("cpu-samples"),       .unit = VALUE_STRING("count")}
 #define CPU_SAMPLES_VALUE_ID 1
-#define WALL_TIME_VALUE         {.type_ = VALUE_STRING("wall-time"),         .unit = VALUE_STRING("nanoseconds")}
 #define WALL_TIME_VALUE_ID 2
-#define ALLOC_SAMPLES_VALUE     {.type_ = VALUE_STRING("alloc-samples"),     .unit = VALUE_STRING("count")}
 #define ALLOC_SAMPLES_VALUE_ID 3
-#define ALLOC_SAMPLES_UNSCALED_VALUE {.type_ = VALUE_STRING("alloc-samples-unscaled"), .unit = VALUE_STRING("count")}
 #define ALLOC_SAMPLES_UNSCALED_VALUE_ID 4
-#define HEAP_SAMPLES_VALUE      {.type_ = VALUE_STRING("heap-live-samples"), .unit = VALUE_STRING("count")}
 #define HEAP_SAMPLES_VALUE_ID 5
-#define HEAP_SIZE_VALUE         {.type_ = VALUE_STRING("heap-live-size"),    .unit = VALUE_STRING("bytes")}
 #define HEAP_SIZE_VALUE_ID 6
-#define TIMELINE_VALUE          {.type_ = VALUE_STRING("timeline"),          .unit = VALUE_STRING("nanoseconds")}
 #define TIMELINE_VALUE_ID 7
 
-static const ddog_prof_ValueType all_value_types[] =
-  {CPU_TIME_VALUE, CPU_SAMPLES_VALUE, WALL_TIME_VALUE, ALLOC_SAMPLES_VALUE, ALLOC_SAMPLES_UNSCALED_VALUE, HEAP_SAMPLES_VALUE, HEAP_SIZE_VALUE, TIMELINE_VALUE};
+static const ddog_prof_SampleType all_sample_types[] = {
+  DDOG_PROF_SAMPLE_TYPE_CPU_TIME,
+  DDOG_PROF_SAMPLE_TYPE_CPU_SAMPLES,
+  DDOG_PROF_SAMPLE_TYPE_WALL_TIME,
+  DDOG_PROF_SAMPLE_TYPE_ALLOC_SAMPLES,
+  DDOG_PROF_SAMPLE_TYPE_ALLOC_SAMPLES_UNSCALED,
+  DDOG_PROF_SAMPLE_TYPE_HEAP_LIVE_SAMPLES,
+  DDOG_PROF_SAMPLE_TYPE_HEAP_LIVE_SIZE,
+  DDOG_PROF_SAMPLE_TYPE_TIMELINE,
+};
 
-// This array MUST be kept in sync with all_value_types above and is intended to act as a "hashmap" between VALUE_ID and the position it
-// occupies on the all_value_types array.
-// E.g. all_value_types_positions[CPU_TIME_VALUE_ID] => 0, means that CPU_TIME_VALUE was declared at position 0 of all_value_types.
+// This array MUST be kept in sync with all_sample_types above and is intended to act as a "hashmap" between VALUE_ID and the position it
+// occupies on the all_sample_types array.
 static const uint8_t all_value_types_positions[] =
   {CPU_TIME_VALUE_ID, CPU_SAMPLES_VALUE_ID, WALL_TIME_VALUE_ID, ALLOC_SAMPLES_VALUE_ID, ALLOC_SAMPLES_UNSCALED_VALUE_ID, HEAP_SAMPLES_VALUE_ID, HEAP_SIZE_VALUE_ID, TIMELINE_VALUE_ID};
 
-#define ALL_VALUE_TYPES_COUNT (sizeof(all_value_types) / sizeof(ddog_prof_ValueType))
+#define ALL_VALUE_TYPES_COUNT (sizeof(all_sample_types) / sizeof(ddog_prof_SampleType))
 
 // Struct for storing stats related to a profile in a particular slot.
 // These stats will share the same lifetime as the data in that profile slot.
-typedef struct slot_stats {
+typedef struct {
   // How many individual samples were recorded into this slot (un-weighted)
   uint64_t recorded_samples;
 } stats_slot;
 
-typedef struct profile_slot {
+typedef struct {
   ddog_prof_Profile profile;
   stats_slot stats;
+  ddog_Timespec start_timestamp;
 } profile_slot;
 
 // Contains native state for each instance
-struct stack_recorder_state {
+typedef struct {
   // Heap recorder instance
   heap_recorder *heap_recorder;
+  bool heap_clean_after_gc_enabled;
+
+  // When set, _native_serialize will call thread_context_collector_on_serialize on this instance
+  // before serializing, so that threads suspended across the whole profile period still get sampled.
+  VALUE thread_context_collector_instance;
 
   pthread_mutex_t mutex_slot_one;
   profile_slot profile_slot_one;
   pthread_mutex_t mutex_slot_two;
   profile_slot profile_slot_two;
 
+  ddog_prof_ManagedStringStorage string_storage;
+  ddog_prof_ManagedStringId label_key_allocation_class;
+  ddog_prof_ManagedStringId label_key_gc_gen_age;
+
   short active_slot; // MUST NEVER BE ACCESSED FROM record_sample; this is NOT for the sampler thread to use.
 
   uint8_t position_for[ALL_VALUE_TYPES_COUNT];
@@ -209,17 +207,17 @@ struct stack_recorder_state {
     long serialization_time_ns_max;
     uint64_t serialization_time_ns_total;
   } stats_lifetime;
-};
+} stack_recorder_state;
 
 // Used to group mutex and the corresponding profile slot for easy unlocking after work is done.
-typedef struct locked_profile_slot {
+typedef struct {
   pthread_mutex_t *mutex;
   profile_slot *data;
 } locked_profile_slot;
 
-struct call_serialize_without_gvl_arguments {
+typedef struct {
   // Set by caller
-  struct stack_recorder_state *state;
+  stack_recorder_state *state;
   ddog_Timespec finish_timestamp;
 
   // Set by callee
@@ -227,50 +225,61 @@ struct call_serialize_without_gvl_arguments {
   ddog_prof_Profile_SerializeResult result;
   long heap_profile_build_time_ns;
   long serialize_no_gvl_time_ns;
+  ddog_prof_MaybeError advance_gen_result;
 
   // Set by both
   bool serialize_ran;
-};
+} call_serialize_without_gvl_arguments;
+
+typedef struct {
+  // Set by caller
+  stack_recorder_state *state;
+  locked_profile_slot active_slot;
+  ddog_prof_Slice_Location locations;
+  sample_values values;
+  sample_labels labels;
+
+  // Set by callee
+  ddog_prof_Profile_Result result;
+} record_sample_arguments;
 
 static VALUE _native_new(VALUE klass);
-static void initialize_slot_concurrency_control(struct stack_recorder_state *state);
-static void initialize_profiles(struct stack_recorder_state *state, ddog_prof_Slice_ValueType sample_types);
+static void initialize_slot_concurrency_control(stack_recorder_state *state);
+static void stack_recorder_typed_data_mark(void *data);
+static void initialize_profiles(stack_recorder_state *state, ddog_prof_Slice_SampleType sample_types);
 static void stack_recorder_typed_data_free(void *data);
-static VALUE _native_initialize(
-  DDTRACE_UNUSED VALUE _self,
-  VALUE recorder_instance,
-  VALUE cpu_time_enabled,
-  VALUE alloc_samples_enabled,
-  VALUE heap_samples_enabled,
-  VALUE heap_size_enabled,
-  VALUE heap_sample_every,
-  VALUE timeline_enabled
-);
+static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _self);
 static VALUE _native_serialize(VALUE self, VALUE recorder_instance);
 static VALUE ruby_time_from(ddog_Timespec ddprof_time);
 static void *call_serialize_without_gvl(void *call_args);
-static locked_profile_slot sampler_lock_active_profile(struct stack_recorder_state *state);
+static VALUE record_sample_locked(VALUE record_sample_arguments_as_value);
+static VALUE record_sample_unlock_ensure(VALUE record_sample_arguments_as_value);
+static locked_profile_slot sampler_lock_active_profile(stack_recorder_state *state);
 static void sampler_unlock_active_profile(locked_profile_slot active_slot);
-static profile_slot* serializer_flip_active_and_inactive_slots(struct stack_recorder_state *state);
+static profile_slot* serializer_flip_active_and_inactive_slots(stack_recorder_state *state);
 static VALUE _native_active_slot(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE _native_is_slot_one_mutex_locked(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE _native_is_slot_two_mutex_locked(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE test_slot_mutex_state(VALUE recorder_instance, int slot);
 static ddog_Timespec system_epoch_now_timespec(void);
 static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE recorder_instance);
-static void serializer_set_start_timestamp_for_next_profile(struct stack_recorder_state *state, ddog_Timespec start_time);
+static void serializer_set_start_timestamp_for_next_profile(stack_recorder_state *state, ddog_Timespec start_time);
 static VALUE _native_record_endpoint(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE local_root_span_id, VALUE endpoint);
-static void reset_profile_slot(profile_slot *slot, ddog_Timespec *start_time /* Can be null */);
-static VALUE _native_track_object(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE new_obj, VALUE weight, VALUE alloc_class);
-static VALUE _native_check_heap_hashes(DDTRACE_UNUSED VALUE _self, VALUE locations);
+static void reset_profile_slot(profile_slot *slot, ddog_Timespec start_timestamp);
 static VALUE _native_start_fake_slow_heap_serialization(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE _native_end_fake_slow_heap_serialization(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE _native_debug_heap_recorder(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
-static VALUE _native_gc_force_recycle(DDTRACE_UNUSED VALUE _self, VALUE obj);
-static VALUE _native_has_seen_id_flag(DDTRACE_UNUSED VALUE _self, VALUE obj);
 static VALUE _native_stats(DDTRACE_UNUSED VALUE self, VALUE instance);
 static VALUE build_profile_stats(profile_slot *slot, long serialization_time_ns, long heap_iteration_prep_time_ns, long heap_profile_build_time_ns);
-
+static VALUE _native_is_object_recorded(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE record_id);
+static VALUE _native_record_id_for(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE obj);
+static VALUE _native_heap_recorder_reset_last_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
+static VALUE _native_heap_recorder_exhaust_record_ids(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
+static VALUE _native_heap_recorder_set_object_alloc_gen(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE record_id, VALUE alloc_gen);
+static VALUE _native_recorder_heap_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
+static VALUE _native_benchmark_intern(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE string, VALUE times, VALUE use_all);
+static VALUE _native_test_managed_string_storage_produces_valid_profiles(DDTRACE_UNUSED VALUE _self);
+static VALUE _native_commit_heap_recordings(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 
 void stack_recorder_init(VALUE profiling_module) {
   VALUE stack_recorder_class = rb_define_class_under(profiling_module, "StackRecorder", rb_cObject);
@@ -287,7 +296,7 @@ void stack_recorder_init(VALUE profiling_module) {
   // https://bugs.ruby-lang.org/issues/18007 for a discussion around this.
   rb_define_alloc_func(stack_recorder_class, _native_new);
 
-  rb_define_singleton_method(stack_recorder_class, "_native_initialize", _native_initialize, 7);
+  rb_define_singleton_method(stack_recorder_class, "_native_initialize", _native_initialize, -1);
   rb_define_singleton_method(stack_recorder_class, "_native_serialize",  _native_serialize, 1);
   rb_define_singleton_method(stack_recorder_class, "_native_reset_after_fork", _native_reset_after_fork, 1);
   rb_define_singleton_method(stack_recorder_class, "_native_stats", _native_stats, 1);
@@ -295,18 +304,21 @@ void stack_recorder_init(VALUE profiling_module) {
   rb_define_singleton_method(testing_module, "_native_slot_one_mutex_locked?", _native_is_slot_one_mutex_locked, 1);
   rb_define_singleton_method(testing_module, "_native_slot_two_mutex_locked?", _native_is_slot_two_mutex_locked, 1);
   rb_define_singleton_method(testing_module, "_native_record_endpoint", _native_record_endpoint, 3);
-  rb_define_singleton_method(testing_module, "_native_track_object", _native_track_object, 4);
-  rb_define_singleton_method(testing_module, "_native_check_heap_hashes", _native_check_heap_hashes, 1);
   rb_define_singleton_method(testing_module, "_native_start_fake_slow_heap_serialization",
       _native_start_fake_slow_heap_serialization, 1);
   rb_define_singleton_method(testing_module, "_native_end_fake_slow_heap_serialization",
       _native_end_fake_slow_heap_serialization, 1);
   rb_define_singleton_method(testing_module, "_native_debug_heap_recorder",
       _native_debug_heap_recorder, 1);
-  rb_define_singleton_method(testing_module, "_native_gc_force_recycle",
-      _native_gc_force_recycle, 1);
-  rb_define_singleton_method(testing_module, "_native_has_seen_id_flag",
-      _native_has_seen_id_flag, 1);
+  rb_define_singleton_method(testing_module, "_native_is_object_recorded?", _native_is_object_recorded, 2);
+  rb_define_singleton_method(testing_module, "_native_record_id_for", _native_record_id_for, 2);
+  rb_define_singleton_method(testing_module, "_native_heap_recorder_reset_last_update", _native_heap_recorder_reset_last_update, 1);
+  rb_define_singleton_method(testing_module, "_native_heap_recorder_exhaust_record_ids", _native_heap_recorder_exhaust_record_ids, 1);
+  rb_define_singleton_method(testing_module, "_native_heap_recorder_set_object_alloc_gen", _native_heap_recorder_set_object_alloc_gen, 3);
+  rb_define_singleton_method(testing_module, "_native_recorder_heap_update", _native_recorder_heap_update, 1);
+  rb_define_singleton_method(testing_module, "_native_benchmark_intern", _native_benchmark_intern, 4);
+  rb_define_singleton_method(testing_module, "_native_test_managed_string_storage_produces_valid_profiles", _native_test_managed_string_storage_produces_valid_profiles, 0);
+  rb_define_singleton_method(testing_module, "_native_commit_heap_recordings", _native_commit_heap_recordings, 1);
 
   ok_symbol = ID2SYM(rb_intern_const("ok"));
   error_symbol = ID2SYM(rb_intern_const("error"));
@@ -317,20 +329,23 @@ void stack_recorder_init(VALUE profiling_module) {
 static const rb_data_type_t stack_recorder_typed_data = {
   .wrap_struct_name = "Datadog::Profiling::StackRecorder",
   .function = {
+    .dmark = stack_recorder_typed_data_mark,
     .dfree = stack_recorder_typed_data_free,
     .dsize = NULL, // We don't track profile memory usage (although it'd be cool if we did!)
-    // No need to provide dmark nor dcompact because we don't directly reference Ruby VALUEs from inside this object
   },
   .flags = RUBY_TYPED_FREE_IMMEDIATELY
 };
 
 static VALUE _native_new(VALUE klass) {
-  struct stack_recorder_state *state = ruby_xcalloc(1, sizeof(struct stack_recorder_state));
+  stack_recorder_state *state = ruby_xcalloc(1, sizeof(stack_recorder_state));
 
   // Note: Any exceptions raised from this note until the TypedData_Wrap_Struct call will lead to the state memory
   // being leaked.
 
-  ddog_prof_Slice_ValueType sample_types = {.ptr = all_value_types, .len = ALL_VALUE_TYPES_COUNT};
+  state->heap_clean_after_gc_enabled = false;
+  state->thread_context_collector_instance = Qnil;
+
+  ddog_prof_Slice_SampleType sample_types = {.ptr = all_sample_types, .len = ALL_VALUE_TYPES_COUNT};
 
   initialize_slot_concurrency_control(state);
   for (uint8_t i = 0; i < ALL_VALUE_TYPES_COUNT; i++) { state->position_for[i] = all_value_types_positions[i]; }
@@ -339,29 +354,32 @@ static VALUE _native_new(VALUE klass) {
     .serialization_time_ns_min = INT64_MAX,
   };
 
-  // Note: At this point, slot_one_profile and slot_two_profile contain null pointers. Libdatadog validates pointers
+  // Note: At this point, slot_one_profile/slot_two_profile/string_storage contain null pointers. Libdatadog validates pointers
   // before using them so it's ok for us to go ahead and create the StackRecorder object.
 
-  // Note: As of this writing, no new Ruby objects get created and stored in the state. If that ever changes, remember
-  // to keep them on the stack and mark them with RB_GC_GUARD -- otherwise it's possible for a GC to run and
-  // since the instance representing the state does not yet exist, such objects will not get marked.
-
   VALUE stack_recorder = TypedData_Wrap_Struct(klass, &stack_recorder_typed_data, state);
 
-  // NOTE: We initialize this because we want a new recorder to be operational even without initialization and our
-  //       default is everything enabled. However, if during recording initialization it turns out we don't want
-  //       heap samples, we will free and reset heap_recorder to NULL, effectively disabling all behaviour specific
-  //       to heap profiling (all calls to heap_recorder_* with a NULL heap recorder are noops).
-  state->heap_recorder = heap_recorder_new();
+  ddog_prof_ManagedStringStorageNewResult string_storage = ddog_prof_ManagedStringStorage_new();
 
-  // Note: Don't raise exceptions after this point, since it'll lead to libdatadog memory leaking!
+  if (string_storage.tag == DDOG_PROF_MANAGED_STRING_STORAGE_NEW_RESULT_ERR) {
+    raise_error(rb_eRuntimeError, "Failed to initialize string storage: %"PRIsVALUE, get_error_details_and_drop(&string_storage.err));
+  }
+
+  state->string_storage = string_storage.ok;
+  state->label_key_allocation_class = intern_or_raise(state->string_storage, DDOG_CHARSLICE_C("allocation class"));
+  state->label_key_gc_gen_age = intern_or_raise(state->string_storage, DDOG_CHARSLICE_C("gc gen age"));
 
   initialize_profiles(state, sample_types);
 
+  // NOTE: We initialize this because we want a new recorder to be operational even before #initialize runs and our
+  //       default is everything enabled. However, if during recording initialization it turns out we don't want
+  //       heap samples, we will free and reset heap_recorder back to NULL.
+  state->heap_recorder = heap_recorder_new(state->string_storage);
+
   return stack_recorder;
 }
 
-static void initialize_slot_concurrency_control(struct stack_recorder_state *state) {
+static void initialize_slot_concurrency_control(stack_recorder_state *state) {
   state->mutex_slot_one = (pthread_mutex_t) PTHREAD_MUTEX_INITIALIZER;
   state->mutex_slot_two = (pthread_mutex_t) PTHREAD_MUTEX_INITIALIZER;
 
@@ -371,34 +389,38 @@ static void initialize_slot_concurrency_control(struct stack_recorder_state *sta
   state->active_slot = 1;
 }
 
-static void initialize_profiles(struct stack_recorder_state *state, ddog_prof_Slice_ValueType sample_types) {
+static void initialize_profiles(stack_recorder_state *state, ddog_prof_Slice_SampleType sample_types) {
+  ddog_Timespec start_timestamp = system_epoch_now_timespec();
+
   ddog_prof_Profile_NewResult slot_one_profile_result =
-    ddog_prof_Profile_new(sample_types, NULL /* period is optional */, NULL /* start_time is optional */);
+    ddog_prof_Profile_with_string_storage(sample_types, NULL /* period is optional */, state->string_storage);
 
   if (slot_one_profile_result.tag == DDOG_PROF_PROFILE_NEW_RESULT_ERR) {
-    rb_raise(rb_eRuntimeError, "Failed to initialize slot one profile: %"PRIsVALUE, get_error_details_and_drop(&slot_one_profile_result.err));
+    raise_error(rb_eRuntimeError, "Failed to initialize slot one profile: %"PRIsVALUE, get_error_details_and_drop(&slot_one_profile_result.err));
   }
 
+  state->profile_slot_one = (profile_slot) { .profile = slot_one_profile_result.ok, .start_timestamp = start_timestamp };
+
   ddog_prof_Profile_NewResult slot_two_profile_result =
-    ddog_prof_Profile_new(sample_types, NULL /* period is optional */, NULL /* start_time is optional */);
+    ddog_prof_Profile_with_string_storage(sample_types, NULL /* period is optional */, state->string_storage);
 
   if (slot_two_profile_result.tag == DDOG_PROF_PROFILE_NEW_RESULT_ERR) {
-    // Uff! Though spot. We need to make sure to properly clean up the other profile as well first
-    ddog_prof_Profile_drop(&slot_one_profile_result.ok);
-    // And now we can raise...
-    rb_raise(rb_eRuntimeError, "Failed to initialize slot two profile: %"PRIsVALUE, get_error_details_and_drop(&slot_two_profile_result.err));
+    // Note: No need to take any special care of slot one, it'll get cleaned up by stack_recorder_typed_data_free
+    raise_error(rb_eRuntimeError, "Failed to initialize slot two profile: %"PRIsVALUE, get_error_details_and_drop(&slot_two_profile_result.err));
   }
 
-  state->profile_slot_one = (profile_slot) {
-    .profile = slot_one_profile_result.ok,
-  };
-  state->profile_slot_two = (profile_slot) {
-    .profile = slot_two_profile_result.ok,
-  };
+  state->profile_slot_two = (profile_slot) { .profile = slot_two_profile_result.ok, .start_timestamp = start_timestamp };
+}
+
+static void stack_recorder_typed_data_mark(void *state_ptr) {
+  stack_recorder_state *state = (stack_recorder_state *) state_ptr;
+
+  rb_gc_mark(state->thread_context_collector_instance);
+  heap_recorder_mark(state->heap_recorder);
 }
 
 static void stack_recorder_typed_data_free(void *state_ptr) {
-  struct stack_recorder_state *state = (struct stack_recorder_state *) state_ptr;
+  stack_recorder_state *state = (stack_recorder_state *) state_ptr;
 
   pthread_mutex_destroy(&state->mutex_slot_one);
   ddog_prof_Profile_drop(&state->profile_slot_one.profile);
@@ -408,37 +430,40 @@ static void stack_recorder_typed_data_free(void *state_ptr) {
 
   heap_recorder_free(state->heap_recorder);
 
+  ddog_prof_ManagedStringStorage_drop(state->string_storage);
+
   ruby_xfree(state);
 }
 
-static VALUE _native_initialize(
-  DDTRACE_UNUSED VALUE _self,
-  VALUE recorder_instance,
-  VALUE cpu_time_enabled,
-  VALUE alloc_samples_enabled,
-  VALUE heap_samples_enabled,
-  VALUE heap_size_enabled,
-  VALUE heap_sample_every,
-  VALUE timeline_enabled
-) {
-  ENFORCE_BOOLEAN(cpu_time_enabled);
+static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _self) {
+  VALUE options;
+  rb_scan_args(argc, argv, "0:", &options);
+  if (options == Qnil) options = rb_hash_new();
+
+  VALUE recorder_instance = rb_hash_fetch(options, ID2SYM(rb_intern("self_instance")));
+  VALUE alloc_samples_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("alloc_samples_enabled")));
+  VALUE heap_samples_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_samples_enabled")));
+  VALUE heap_size_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_size_enabled")));
+  VALUE heap_sample_every = rb_hash_fetch(options, ID2SYM(rb_intern("heap_sample_every")));
+  VALUE heap_clean_after_gc_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_clean_after_gc_enabled")));
+
   ENFORCE_BOOLEAN(alloc_samples_enabled);
   ENFORCE_BOOLEAN(heap_samples_enabled);
   ENFORCE_BOOLEAN(heap_size_enabled);
   ENFORCE_TYPE(heap_sample_every, T_FIXNUM);
-  ENFORCE_BOOLEAN(timeline_enabled);
+  ENFORCE_BOOLEAN(heap_clean_after_gc_enabled);
 
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  state->heap_clean_after_gc_enabled = (heap_clean_after_gc_enabled == Qtrue);
 
   heap_recorder_set_sample_rate(state->heap_recorder, NUM2INT(heap_sample_every));
 
   uint8_t requested_values_count = ALL_VALUE_TYPES_COUNT -
-    (cpu_time_enabled == Qtrue ? 0 : 1) -
     (alloc_samples_enabled == Qtrue? 0 : 2) -
     (heap_samples_enabled == Qtrue ? 0 : 1) -
-    (heap_size_enabled == Qtrue ? 0 : 1) -
-    (timeline_enabled == Qtrue ? 0 : 1);
+    (heap_size_enabled == Qtrue ? 0 : 1);
 
   if (requested_values_count == ALL_VALUE_TYPES_COUNT) return Qtrue; // Nothing to do, this is the default
 
@@ -448,30 +473,31 @@ static VALUE _native_initialize(
 
   state->enabled_values_count = requested_values_count;
 
-  ddog_prof_ValueType enabled_value_types[ALL_VALUE_TYPES_COUNT];
+  ddog_prof_SampleType enabled_sample_types[ALL_VALUE_TYPES_COUNT];
   uint8_t next_enabled_pos = 0;
   uint8_t next_disabled_pos = requested_values_count;
 
-  // CPU_SAMPLES_VALUE is always enabled
-  enabled_value_types[next_enabled_pos] = (ddog_prof_ValueType) CPU_SAMPLES_VALUE;
+  // CPU_SAMPLES is always enabled
+  enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_CPU_SAMPLES;
   state->position_for[CPU_SAMPLES_VALUE_ID] = next_enabled_pos++;
 
-  // WALL_TIME_VALUE is always enabled
-  enabled_value_types[next_enabled_pos] = (ddog_prof_ValueType) WALL_TIME_VALUE;
+  // WALL_TIME is always enabled
+  enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_WALL_TIME;
   state->position_for[WALL_TIME_VALUE_ID] = next_enabled_pos++;
 
-  if (cpu_time_enabled == Qtrue) {
-    enabled_value_types[next_enabled_pos] = (ddog_prof_ValueType) CPU_TIME_VALUE;
-    state->position_for[CPU_TIME_VALUE_ID] = next_enabled_pos++;
-  } else {
-    state->position_for[CPU_TIME_VALUE_ID] = next_disabled_pos++;
-  }
+  // CPU_TIME is always enabled
+  enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_CPU_TIME;
+  state->position_for[CPU_TIME_VALUE_ID] = next_enabled_pos++;
+
+  // TIMELINE is always enabled
+  enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_TIMELINE;
+  state->position_for[TIMELINE_VALUE_ID] = next_enabled_pos++;
 
   if (alloc_samples_enabled == Qtrue) {
-    enabled_value_types[next_enabled_pos] = (ddog_prof_ValueType) ALLOC_SAMPLES_VALUE;
+    enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_ALLOC_SAMPLES;
     state->position_for[ALLOC_SAMPLES_VALUE_ID] = next_enabled_pos++;
 
-    enabled_value_types[next_enabled_pos] = (ddog_prof_ValueType) ALLOC_SAMPLES_UNSCALED_VALUE;
+    enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_ALLOC_SAMPLES_UNSCALED;
     state->position_for[ALLOC_SAMPLES_UNSCALED_VALUE_ID] = next_enabled_pos++;
   } else {
     state->position_for[ALLOC_SAMPLES_VALUE_ID] = next_disabled_pos++;
@@ -479,14 +505,14 @@ static VALUE _native_initialize(
   }
 
   if (heap_samples_enabled == Qtrue) {
-    enabled_value_types[next_enabled_pos] = (ddog_prof_ValueType) HEAP_SAMPLES_VALUE;
+    enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_HEAP_LIVE_SAMPLES;
     state->position_for[HEAP_SAMPLES_VALUE_ID] = next_enabled_pos++;
   } else {
     state->position_for[HEAP_SAMPLES_VALUE_ID] = next_disabled_pos++;
   }
 
   if (heap_size_enabled == Qtrue) {
-    enabled_value_types[next_enabled_pos] = (ddog_prof_ValueType) HEAP_SIZE_VALUE;
+    enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_HEAP_LIVE_SIZE;
     state->position_for[HEAP_SIZE_VALUE_ID] = next_enabled_pos++;
   } else {
     state->position_for[HEAP_SIZE_VALUE_ID] = next_disabled_pos++;
@@ -500,42 +526,41 @@ static VALUE _native_initialize(
     state->heap_recorder = NULL;
   }
 
-  if (timeline_enabled == Qtrue) {
-    enabled_value_types[next_enabled_pos] = (ddog_prof_ValueType) TIMELINE_VALUE;
-    state->position_for[TIMELINE_VALUE_ID] = next_enabled_pos++;
-  } else {
-    state->position_for[TIMELINE_VALUE_ID] = next_disabled_pos++;
-  }
-
   ddog_prof_Profile_drop(&state->profile_slot_one.profile);
   ddog_prof_Profile_drop(&state->profile_slot_two.profile);
 
-  ddog_prof_Slice_ValueType sample_types = {.ptr = enabled_value_types, .len = state->enabled_values_count};
+  ddog_prof_Slice_SampleType sample_types = {.ptr = enabled_sample_types, .len = state->enabled_values_count};
   initialize_profiles(state, sample_types);
 
   return Qtrue;
 }
 
 static VALUE _native_serialize(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
-  ddog_Timespec finish_timestamp = system_epoch_now_timespec();
-  // Need to do this while still holding on to the Global VM Lock; see comments on method for why
-  serializer_set_start_timestamp_for_next_profile(state, finish_timestamp);
+  if (state->thread_context_collector_instance != Qnil) {
+    thread_context_collector_on_serialize(state->thread_context_collector_instance);
+  }
 
   long heap_iteration_prep_start_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
   // Prepare the iteration on heap recorder we'll be doing outside the GVL. The preparation needs to
-  // happen while holding on to the GVL.
+  // happen while holding the GVL.
+  // NOTE: While rare, it's possible for the GVL to be released inside this function (see comments on `heap_recorder_update`)
+  // and thus don't assume this is an "atomic" step -- other threads may get some running time in the meanwhile.
   heap_recorder_prepare_iteration(state->heap_recorder);
   long heap_iteration_prep_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE) - heap_iteration_prep_start_time_ns;
 
+  ddog_Timespec finish_timestamp = system_epoch_now_timespec();
+  // Need to do this while still holding the Global VM Lock; see comments on method for why
+  serializer_set_start_timestamp_for_next_profile(state, finish_timestamp);
+
   // We'll release the Global VM Lock while we're calling serialize, so that the Ruby VM can continue to work while this
   // is pending
-  struct call_serialize_without_gvl_arguments args = {
+  call_serialize_without_gvl_arguments args = {
     .state = state,
     .finish_timestamp = finish_timestamp,
-    .serialize_ran = false
+    .serialize_ran = false,
   };
 
   while (!args.serialize_ran) {
@@ -552,20 +577,16 @@ static VALUE _native_serialize(DDTRACE_UNUSED VALUE _self, VALUE recorder_instan
     rb_thread_call_without_gvl2(call_serialize_without_gvl, &args, NULL /* No interruption function needed in this case */, NULL /* Not needed */);
   }
 
-  // Cleanup after heap recorder iteration. This needs to happen while holding on to the GVL.
+  // Cleanup after heap recorder iteration. This needs to happen while holding the GVL.
   heap_recorder_finish_iteration(state->heap_recorder);
 
   // NOTE: We are focusing on the serialization time outside of the GVL in this stat here. This doesn't
   //       really cover the full serialization process but it gives a more useful number since it bypasses
   //       the noise of acquiring GVLs and dealing with interruptions which is highly specific to runtime
   //       conditions and over which we really have no control about.
-  long serialization_time_ns = args.serialize_no_gvl_time_ns;
-  if (serialization_time_ns >= 0) {
-    // Only update stats if our serialization time is valid.
-    state->stats_lifetime.serialization_time_ns_max = long_max_of(state->stats_lifetime.serialization_time_ns_max, serialization_time_ns);
-    state->stats_lifetime.serialization_time_ns_min = long_min_of(state->stats_lifetime.serialization_time_ns_min, serialization_time_ns);
-    state->stats_lifetime.serialization_time_ns_total += serialization_time_ns;
-  }
+  state->stats_lifetime.serialization_time_ns_max = long_max_of(state->stats_lifetime.serialization_time_ns_max, args.serialize_no_gvl_time_ns);
+  state->stats_lifetime.serialization_time_ns_min = long_min_of(state->stats_lifetime.serialization_time_ns_min, args.serialize_no_gvl_time_ns);
+  state->stats_lifetime.serialization_time_ns_total += args.serialize_no_gvl_time_ns;
 
   ddog_prof_Profile_SerializeResult serialized_profile = args.result;
 
@@ -574,20 +595,22 @@ static VALUE _native_serialize(DDTRACE_UNUSED VALUE _self, VALUE recorder_instan
     return rb_ary_new_from_args(2, error_symbol, get_error_details_and_drop(&serialized_profile.err));
   }
 
+  // Note: If we got here, the profile serialized correctly.
+  // Once we wrap this into a Ruby object, our `EncodedProfile` class will automatically manage memory for it and we
+  // can raise exceptions without worrying about leaking the profile.
   state->stats_lifetime.serialization_successes++;
+  VALUE encoded_profile = from_ddog_prof_EncodedProfile(serialized_profile.ok);
 
-  VALUE encoded_pprof = ruby_string_from_vec_u8(serialized_profile.ok.buffer);
-
-  ddog_Timespec ddprof_start = serialized_profile.ok.start;
-  ddog_Timespec ddprof_finish = serialized_profile.ok.end;
-
-  ddog_prof_EncodedProfile_drop(&serialized_profile.ok);
+  ddog_prof_MaybeError result = args.advance_gen_result;
+  if (result.tag == DDOG_PROF_OPTION_ERROR_SOME_ERROR) {
+    raise_error(rb_eRuntimeError, "Failed to advance string storage gen: %"PRIsVALUE, get_error_details_and_drop(&result.some));
+  }
 
-  VALUE start = ruby_time_from(ddprof_start);
-  VALUE finish = ruby_time_from(ddprof_finish);
-  VALUE profile_stats = build_profile_stats(args.slot, serialization_time_ns, heap_iteration_prep_time_ns, args.heap_profile_build_time_ns);
+  VALUE start = ruby_time_from(args.slot->start_timestamp);
+  VALUE finish = ruby_time_from(finish_timestamp);
+  VALUE profile_stats = build_profile_stats(args.slot, args.serialize_no_gvl_time_ns, heap_iteration_prep_time_ns, args.heap_profile_build_time_ns);
 
-  return rb_ary_new_from_args(2, ok_symbol, rb_ary_new_from_args(4, start, finish, encoded_pprof, profile_stats));
+  return rb_ary_new_from_args(2, ok_symbol, rb_ary_new_from_args(4, start, finish, encoded_profile, profile_stats));
 }
 
 static VALUE ruby_time_from(ddog_Timespec ddprof_time) {
@@ -597,10 +620,31 @@ static VALUE ruby_time_from(ddog_Timespec ddprof_time) {
 }
 
 void record_sample(VALUE recorder_instance, ddog_prof_Slice_Location locations, sample_values values, sample_labels labels) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
-  locked_profile_slot active_slot = sampler_lock_active_profile(state);
+  record_sample_arguments args = {
+    .state = state,
+    .active_slot = sampler_lock_active_profile(state),
+    .locations = locations,
+    .values = values,
+    .labels = labels,
+  };
+
+  // NOTE: The heap recorder is allowed to raise exceptions if something's wrong. We use `rb_ensure` to make sure we
+  // properly unlock the active slot mutex on our way out. Otherwise, this would later lead to deadlocks (since the
+  // active slot mutex is not expected to be locked forever).
+  rb_ensure(record_sample_locked, (VALUE) &args, record_sample_unlock_ensure, (VALUE) &args);
+
+  if (args.result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
+    raise_error(rb_eArgError, "Failed to record sample: %"PRIsVALUE, get_error_details_and_drop(&args.result.err));
+  }
+}
+
+static VALUE record_sample_locked(VALUE record_sample_arguments_as_value) {
+  record_sample_arguments *args = (record_sample_arguments *) record_sample_arguments_as_value;
+  stack_recorder_state *state = args->state;
+  sample_values values = args->values;
 
   // Note: We initialize this array to have ALL_VALUE_TYPES_COUNT but only tell libdatadog to use the first
   // state->enabled_values_count values. This simplifies handling disabled value types -- we still put them on the
@@ -616,53 +660,44 @@ void record_sample(VALUE recorder_instance, ddog_prof_Slice_Location locations,
   metric_values[position_for[ALLOC_SAMPLES_UNSCALED_VALUE_ID]] = values.alloc_samples_unscaled;
   metric_values[position_for[TIMELINE_VALUE_ID]]      = values.timeline_wall_time_ns;
 
-  if (values.heap_sample) {
-    // If we got an allocation sample end the heap allocation recording to commit the heap sample.
-    // FIXME: Heap sampling currently has to be done in 2 parts because the construction of locations is happening
-    //        very late in the allocation-sampling path (which is shared with the cpu sampling path). This can
-    //        be fixed with some refactoring but for now this leads to a less impactful change.
-    //
-    // NOTE: The heap recorder is allowed to raise exceptions if something's wrong. But we also need to handle it
-    // on this side to make sure we properly unlock the active slot mutex on our way out. Otherwise, this would
-    // later lead to deadlocks (since the active slot mutex is not expected to be locked forever).
-    int exception_state = end_heap_allocation_recording_with_rb_protect(state->heap_recorder, locations);
-    if (exception_state) {
-      sampler_unlock_active_profile(active_slot);
-      rb_jump_tag(exception_state);
-    }
+  if (values.heap_sample != NULL) {
+    heap_recorder_record_allocation(
+      state->heap_recorder,
+      values.heap_sample->new_object,
+      values.alloc_samples,
+      values.heap_sample->alloc_class,
+      args->locations,
+      &values.heap_sample->out_needs_commit
+    );
   }
 
-  ddog_prof_Profile_Result result = ddog_prof_Profile_add(
-    &active_slot.data->profile,
+  args->result = ddog_prof_Profile_add(
+    &args->active_slot.data->profile,
     (ddog_prof_Sample) {
-      .locations = locations,
+      .locations = args->locations,
       .values = (ddog_Slice_I64) {.ptr = metric_values, .len = state->enabled_values_count},
-      .labels = labels.labels
+      .labels = args->labels.labels
     },
-    labels.end_timestamp_ns
+    /* To disable end_timestamp_ns to get aggregated profiles in pprof, replace the below with `0` */
+    args->labels.end_timestamp_ns
   );
 
-  active_slot.data->stats.recorded_samples++;
-
-  sampler_unlock_active_profile(active_slot);
+  args->active_slot.data->stats.recorded_samples++;
 
-  if (result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
-    rb_raise(rb_eArgError, "Failed to record sample: %"PRIsVALUE, get_error_details_and_drop(&result.err));
-  }
+  return Qnil; // Unused
 }
 
-void track_object(VALUE recorder_instance, VALUE new_object, unsigned int sample_weight, ddog_CharSlice *alloc_class) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
-  // FIXME: Heap sampling currently has to be done in 2 parts because the construction of locations is happening
-  //        very late in the allocation-sampling path (which is shared with the cpu sampling path). This can
-  //        be fixed with some refactoring but for now this leads to a less impactful change.
-  start_heap_allocation_recording(state->heap_recorder, new_object, sample_weight, alloc_class);
+static VALUE record_sample_unlock_ensure(VALUE record_sample_arguments_as_value) {
+  record_sample_arguments *args = (record_sample_arguments *) record_sample_arguments_as_value;
+
+  sampler_unlock_active_profile(args->active_slot);
+
+  return Qnil; // Unused
 }
 
 void record_endpoint(VALUE recorder_instance, uint64_t local_root_span_id, ddog_CharSlice endpoint) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   locked_profile_slot active_slot = sampler_lock_active_profile(state);
 
@@ -671,16 +706,30 @@ void record_endpoint(VALUE recorder_instance, uint64_t local_root_span_id, ddog_
   sampler_unlock_active_profile(active_slot);
 
   if (result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
-    rb_raise(rb_eArgError, "Failed to record endpoint: %"PRIsVALUE, get_error_details_and_drop(&result.err));
+    raise_error(rb_eArgError, "Failed to record endpoint: %"PRIsVALUE, get_error_details_and_drop(&result.err));
   }
 }
 
+void recorder_heap_update_may_lose_gvl(VALUE recorder_instance) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  if (state->heap_clean_after_gc_enabled) heap_recorder_update_young_objects(state->heap_recorder);
+}
+
+void recorder_commit_heap_recordings_may_lose_gvl(VALUE recorder_instance) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  heap_recorder_commit_recordings_may_lose_gvl(state->heap_recorder);
+}
+
 #define MAX_LEN_HEAP_ITERATION_ERROR_MSG 256
 
 // Heap recorder iteration context allows us access to stack recorder state and profile being serialized
 // during iteration of heap recorder live objects.
 typedef struct heap_recorder_iteration_context {
-  struct stack_recorder_state *state;
+  stack_recorder_state *state;
   profile_slot *slot;
 
   bool error;
@@ -701,18 +750,15 @@ static bool add_heap_sample_to_active_profile_without_gvl(heap_recorder_iteratio
   ddog_prof_Label labels[2];
   size_t label_offset = 0;
 
-  if (object_data->class != NULL) {
+  if (object_data->class.value > 0) {
     labels[label_offset++] = (ddog_prof_Label) {
-      .key = DDOG_CHARSLICE_C("allocation class"),
-      .str = (ddog_CharSlice) {
-        .ptr = object_data->class,
-        .len = strlen(object_data->class),
-      },
+      .key_id = context->state->label_key_allocation_class,
+      .str_id = object_data->class,
       .num = 0, // This shouldn't be needed but the tracer-2.7 docker image ships a buggy gcc that complains about this
     };
   }
   labels[label_offset++] = (ddog_prof_Label) {
-    .key = DDOG_CHARSLICE_C("gc gen age"),
+    .key_id = context->state->label_key_gc_gen_age,
     .num = object_data->gen_age,
   };
 
@@ -742,7 +788,7 @@ static bool add_heap_sample_to_active_profile_without_gvl(heap_recorder_iteratio
   return true;
 }
 
-static void build_heap_profile_without_gvl(struct stack_recorder_state *state, profile_slot *slot) {
+static void build_heap_profile_without_gvl(stack_recorder_state *state, profile_slot *slot) {
   heap_recorder_iteration_context iteration_context = {
     .state = state,
     .slot = slot,
@@ -763,12 +809,11 @@ static void build_heap_profile_without_gvl(struct stack_recorder_state *state, p
 }
 
 static void *call_serialize_without_gvl(void *call_args) {
-  struct call_serialize_without_gvl_arguments *args = (struct call_serialize_without_gvl_arguments *) call_args;
+  call_serialize_without_gvl_arguments *args = (call_serialize_without_gvl_arguments *) call_args;
 
   long serialize_no_gvl_start_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
 
   profile_slot *slot_now_inactive = serializer_flip_active_and_inactive_slots(args->state);
-
   args->slot = slot_now_inactive;
 
   // Now that we have the inactive profile with all but heap samples, lets fill it with heap data
@@ -777,19 +822,20 @@ static void *call_serialize_without_gvl(void *call_args) {
   args->heap_profile_build_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE) - serialize_no_gvl_start_time_ns;
 
   // Note: The profile gets reset by the serialize call
-  args->result = ddog_prof_Profile_serialize(&args->slot->profile, &args->finish_timestamp, NULL /* duration_nanos is optional */, NULL /* start_time is optional */);
+  args->result = ddog_prof_Profile_serialize(&args->slot->profile, &args->slot->start_timestamp, &args->finish_timestamp);
+  args->advance_gen_result = ddog_prof_ManagedStringStorage_advance_gen(args->state->string_storage);
   args->serialize_ran = true;
-  args->serialize_no_gvl_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE) - serialize_no_gvl_start_time_ns;
+  args->serialize_no_gvl_time_ns = long_max_of(0, monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE) - serialize_no_gvl_start_time_ns);
 
   return NULL; // Unused
 }
 
 VALUE enforce_recorder_instance(VALUE object) {
-  Check_TypedStruct(object, &stack_recorder_typed_data);
+  ENFORCE_TYPED_DATA(object, &stack_recorder_typed_data);
   return object;
 }
 
-static locked_profile_slot sampler_lock_active_profile(struct stack_recorder_state *state) {
+static locked_profile_slot sampler_lock_active_profile(stack_recorder_state *state) {
   int error;
 
   for (int attempts = 0; attempts < 2; attempts++) {
@@ -809,14 +855,14 @@ static locked_profile_slot sampler_lock_active_profile(struct stack_recorder_sta
   }
 
   // We already tried both multiple times, and we did not succeed. This is not expected to happen. Let's stop sampling.
-  rb_raise(rb_eRuntimeError, "Failed to grab either mutex in sampler_lock_active_profile");
+  raise_error(rb_eRuntimeError, "Failed to grab either mutex in sampler_lock_active_profile");
 }
 
 static void sampler_unlock_active_profile(locked_profile_slot active_slot) {
   ENFORCE_SUCCESS_GVL(pthread_mutex_unlock(active_slot.mutex));
 }
 
-static profile_slot* serializer_flip_active_and_inactive_slots(struct stack_recorder_state *state) {
+static profile_slot* serializer_flip_active_and_inactive_slots(stack_recorder_state *state) {
   int previously_active_slot = state->active_slot;
 
   if (previously_active_slot != 1 && previously_active_slot != 2) {
@@ -842,8 +888,8 @@ static profile_slot* serializer_flip_active_and_inactive_slots(struct stack_reco
 // This method exists only to enable testing Datadog::Profiling::StackRecorder behavior using RSpec.
 // It SHOULD NOT be used for other purposes.
 static VALUE _native_active_slot(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   return INT2NUM(state->active_slot);
 }
@@ -857,8 +903,8 @@ static VALUE _native_is_slot_one_mutex_locked(DDTRACE_UNUSED VALUE _self, VALUE
 static VALUE _native_is_slot_two_mutex_locked(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) { return test_slot_mutex_state(recorder_instance, 2); }
 
 static VALUE test_slot_mutex_state(VALUE recorder_instance, int slot) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   pthread_mutex_t *slot_mutex = (slot == 1) ? &state->mutex_slot_one : &state->mutex_slot_two;
 
@@ -874,7 +920,7 @@ static VALUE test_slot_mutex_state(VALUE recorder_instance, int slot) {
     return Qtrue;
   } else {
     ENFORCE_SUCCESS_GVL(error);
-    rb_raise(rb_eRuntimeError, "Failed to raise exception in test_slot_mutex_state; this should never happen");
+    raise_error(rb_eRuntimeError, "Failed to raise exception in test_slot_mutex_state; this should never happen");
   }
 }
 
@@ -887,16 +933,19 @@ static ddog_Timespec system_epoch_now_timespec(void) {
 //
 // Assumption: This method gets called BEFORE restarting profiling -- e.g. there are no components attempting to
 // trigger samples at the same time.
+//
+// Note that tests call this method directly in the same process without forking,
+// and in such a case non-current Threads keep running.
 static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE recorder_instance) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   // In case the fork happened halfway through `serializer_flip_active_and_inactive_slots` execution and the
   // resulting state is inconsistent, we make sure to reset it back to the initial state.
   initialize_slot_concurrency_control(state);
-
-  reset_profile_slot(&state->profile_slot_one, /* start_time: */ NULL);
-  reset_profile_slot(&state->profile_slot_two, /* start_time: */ NULL);
+  ddog_Timespec start_timestamp = system_epoch_now_timespec();
+  reset_profile_slot(&state->profile_slot_one, start_timestamp);
+  reset_profile_slot(&state->profile_slot_two, start_timestamp);
 
   heap_recorder_after_fork(state->heap_recorder);
 
@@ -905,10 +954,10 @@ static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE recorder_
 
 // Assumption 1: This method is called with the GVL being held, because `ddog_prof_Profile_reset` mutates the profile and must
 // not be interrupted part-way through by a VM fork.
-static void serializer_set_start_timestamp_for_next_profile(struct stack_recorder_state *state, ddog_Timespec start_time) {
+static void serializer_set_start_timestamp_for_next_profile(stack_recorder_state *state, ddog_Timespec start_time) {
   // Before making this profile active, we reset it so that it uses the correct start_time for its start
   profile_slot *next_profile_slot = (state->active_slot == 1) ? &state->profile_slot_two : &state->profile_slot_one;
-  reset_profile_slot(next_profile_slot, &start_time);
+  reset_profile_slot(next_profile_slot, start_time);
 }
 
 static VALUE _native_record_endpoint(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE local_root_span_id, VALUE endpoint) {
@@ -917,56 +966,20 @@ static VALUE _native_record_endpoint(DDTRACE_UNUSED VALUE _self, VALUE recorder_
   return Qtrue;
 }
 
-static VALUE _native_track_object(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE new_obj, VALUE weight, VALUE alloc_class) {
-  ENFORCE_TYPE(weight, T_FIXNUM);
-  ddog_CharSlice alloc_class_slice = char_slice_from_ruby_string(alloc_class);
-  track_object(recorder_instance, new_obj, NUM2UINT(weight), &alloc_class_slice);
-  return Qtrue;
-}
-
-static VALUE _native_check_heap_hashes(DDTRACE_UNUSED VALUE _self, VALUE locations) {
-  ENFORCE_TYPE(locations, T_ARRAY);
-  size_t locations_len = rb_array_len(locations);
-  ddog_prof_Location locations_arr[locations_len];
-  for (size_t i = 0; i < locations_len; i++) {
-    VALUE location = rb_ary_entry(locations, i);
-    ENFORCE_TYPE(location, T_ARRAY);
-    VALUE name = rb_ary_entry(location, 0);
-    VALUE filename = rb_ary_entry(location, 1);
-    VALUE line = rb_ary_entry(location, 2);
-    ENFORCE_TYPE(name, T_STRING);
-    ENFORCE_TYPE(filename, T_STRING);
-    ENFORCE_TYPE(line, T_FIXNUM);
-    locations_arr[i] = (ddog_prof_Location) {
-      .line = line,
-        .function = (ddog_prof_Function) {
-          .name = char_slice_from_ruby_string(name),
-          .filename = char_slice_from_ruby_string(filename),
-        }
-    };
-  }
-  ddog_prof_Slice_Location ddog_locations = {
-    .len = locations_len,
-    .ptr = locations_arr,
-  };
-  heap_recorder_testonly_assert_hash_matches(ddog_locations);
-
-  return Qnil;
-}
-
-static void reset_profile_slot(profile_slot *slot, ddog_Timespec *start_time /* Can be null */) {
-  ddog_prof_Profile_Result reset_result = ddog_prof_Profile_reset(&slot->profile, start_time);
+static void reset_profile_slot(profile_slot *slot, ddog_Timespec start_timestamp) {
+  ddog_prof_Profile_Result reset_result = ddog_prof_Profile_reset(&slot->profile);
   if (reset_result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
-    rb_raise(rb_eRuntimeError, "Failed to reset profile: %"PRIsVALUE, get_error_details_and_drop(&reset_result.err));
+    raise_error(rb_eRuntimeError, "Failed to reset profile: %"PRIsVALUE, get_error_details_and_drop(&reset_result.err));
   }
+  slot->start_timestamp = start_timestamp;
   slot->stats = (stats_slot) {};
 }
 
 // This method exists only to enable testing Datadog::Profiling::StackRecorder behavior using RSpec.
 // It SHOULD NOT be used for other purposes.
 static VALUE _native_start_fake_slow_heap_serialization(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   heap_recorder_prepare_iteration(state->heap_recorder);
 
@@ -976,8 +989,8 @@ static VALUE _native_start_fake_slow_heap_serialization(DDTRACE_UNUSED VALUE _se
 // This method exists only to enable testing Datadog::Profiling::StackRecorder behavior using RSpec.
 // It SHOULD NOT be used for other purposes.
 static VALUE _native_end_fake_slow_heap_serialization(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   heap_recorder_finish_iteration(state->heap_recorder);
 
@@ -987,43 +1000,15 @@ static VALUE _native_end_fake_slow_heap_serialization(DDTRACE_UNUSED VALUE _self
 // This method exists only to enable testing Datadog::Profiling::StackRecorder behavior using RSpec.
 // It SHOULD NOT be used for other purposes.
 static VALUE _native_debug_heap_recorder(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   return heap_recorder_testonly_debug(state->heap_recorder);
 }
 
-#pragma GCC diagnostic push
-// rb_gc_force_recycle was deprecated in latest versions of Ruby and is a noop.
-#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
-#pragma GCC diagnostic ignored "-Wunused-parameter"
-// This method exists only to enable testing Datadog::Profiling::StackRecorder behavior using RSpec.
-// It SHOULD NOT be used for other purposes.
-static VALUE _native_gc_force_recycle(DDTRACE_UNUSED VALUE _self, VALUE obj) {
-  #ifdef HAVE_WORKING_RB_GC_FORCE_RECYCLE
-    rb_gc_force_recycle(obj);
-  #endif
-  return Qnil;
-}
-#pragma GCC diagnostic pop
-
-// This method exists only to enable testing Datadog::Profiling::StackRecorder behavior using RSpec.
-// It SHOULD NOT be used for other purposes.
-static VALUE _native_has_seen_id_flag(DDTRACE_UNUSED VALUE _self, VALUE obj) {
-  #ifndef NO_SEEN_OBJ_ID_FLAG
-    if (RB_FL_TEST(obj, RUBY_FL_SEEN_OBJ_ID)) {
-      return Qtrue;
-    } else {
-      return Qfalse;
-    }
-  #else
-    return Qfalse;
-  #endif
-}
-
 static VALUE _native_stats(DDTRACE_UNUSED VALUE self, VALUE recorder_instance) {
-  struct stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   uint64_t total_serializations = state->stats_lifetime.serialization_successes + state->stats_lifetime.serialization_failures;
 
@@ -1057,3 +1042,174 @@ static VALUE build_profile_stats(profile_slot *slot, long serialization_time_ns,
   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_is_object_recorded(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE record_id) {
+  ENFORCE_TYPE(record_id, T_FIXNUM);
+
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  return heap_recorder_testonly_is_object_recorded(state->heap_recorder, FIX2LONG(record_id));
+}
+
+static VALUE _native_record_id_for(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE obj) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  return heap_recorder_testonly_record_id_for(state->heap_recorder, obj);
+}
+
+static VALUE _native_heap_recorder_exhaust_record_ids(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  heap_recorder_testonly_exhaust_record_ids(state->heap_recorder);
+  return Qtrue;
+}
+
+static VALUE _native_heap_recorder_set_object_alloc_gen(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE record_id, VALUE alloc_gen) {
+  ENFORCE_TYPE(record_id, T_FIXNUM);
+  ENFORCE_TYPE(alloc_gen, T_FIXNUM);
+
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  heap_recorder_testonly_set_object_alloc_gen(state->heap_recorder, FIX2LONG(record_id), NUM2SIZET(alloc_gen));
+  return Qtrue;
+}
+
+static VALUE _native_heap_recorder_reset_last_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  heap_recorder_testonly_reset_last_update(state->heap_recorder);
+
+  return Qtrue;
+}
+
+static VALUE _native_recorder_heap_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
+  recorder_heap_update_may_lose_gvl(recorder_instance);
+  return Qtrue;
+}
+
+static VALUE _native_benchmark_intern(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE string, VALUE times, VALUE use_all) {
+  ENFORCE_TYPE(string, T_STRING);
+  ENFORCE_TYPE(times, T_FIXNUM);
+  ENFORCE_BOOLEAN(use_all);
+
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  heap_recorder_testonly_benchmark_intern(state->heap_recorder, char_slice_from_ruby_string(string), FIX2INT(times), use_all == Qtrue);
+
+  return Qtrue;
+}
+
+// See comments in rspec test for details on what we're testing here.
+static VALUE _native_test_managed_string_storage_produces_valid_profiles(DDTRACE_UNUSED VALUE _self) {
+  ddog_prof_ManagedStringStorageNewResult string_storage = ddog_prof_ManagedStringStorage_new();
+
+  if (string_storage.tag == DDOG_PROF_MANAGED_STRING_STORAGE_NEW_RESULT_ERR) {
+    raise_error(rb_eRuntimeError, "Failed to initialize string storage: %"PRIsVALUE, get_error_details_and_drop(&string_storage.err));
+  }
+
+  ddog_prof_Slice_SampleType sample_types = {.ptr = all_sample_types, .len = ALL_VALUE_TYPES_COUNT};
+  ddog_prof_Profile_NewResult profile = ddog_prof_Profile_with_string_storage(sample_types, NULL, string_storage.ok);
+
+  if (profile.tag == DDOG_PROF_PROFILE_NEW_RESULT_ERR) {
+    raise_error(rb_eRuntimeError, "Failed to initialize profile: %"PRIsVALUE, get_error_details_and_drop(&profile.err));
+  }
+
+  ddog_prof_ManagedStringId hello = intern_or_raise(string_storage.ok, DDOG_CHARSLICE_C("hello"));
+  ddog_prof_ManagedStringId world = intern_or_raise(string_storage.ok, DDOG_CHARSLICE_C("world"));
+  ddog_prof_ManagedStringId key   = intern_or_raise(string_storage.ok, DDOG_CHARSLICE_C("key"));
+
+  int64_t metric_values[] = {1, 2, 3, 4, 5, 6, 7, 8};
+  ddog_prof_Label labels[] = {{.key_id = key, .str_id = key}};
+
+  ddog_prof_Location locations[] = {
+    (ddog_prof_Location) {
+      .mapping = {.filename = DDOG_CHARSLICE_C(""), .build_id = DDOG_CHARSLICE_C(""), .build_id_id = {}},
+      .function = {
+        .name = DDOG_CHARSLICE_C(""),
+        .name_id = hello,
+        .filename = DDOG_CHARSLICE_C(""),
+        .filename_id = world,
+      },
+      .line = 1,
+    }
+  };
+
+  ddog_prof_Profile_Result result = ddog_prof_Profile_add(
+    &profile.ok,
+    (ddog_prof_Sample) {
+      .locations = (ddog_prof_Slice_Location) { .ptr = locations, .len = 1},
+      .values = (ddog_Slice_I64) {.ptr = metric_values, .len = 8},
+      .labels = (ddog_prof_Slice_Label) { .ptr = labels, .len = 1 }
+    },
+    0
+  );
+
+  if (result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
+    raise_error(rb_eArgError, "Failed to record sample: %"PRIsVALUE, get_error_details_and_drop(&result.err));
+  }
+
+  ddog_Timespec finish_timestamp = system_epoch_now_timespec();
+  ddog_Timespec start_timestamp = {.seconds = finish_timestamp.seconds - 60};
+  ddog_prof_Profile_SerializeResult serialize_result = ddog_prof_Profile_serialize(&profile.ok, &start_timestamp, &finish_timestamp);
+
+  if (serialize_result.tag == DDOG_PROF_PROFILE_SERIALIZE_RESULT_ERR) {
+    raise_error(rb_eRuntimeError, "Failed to serialize: %"PRIsVALUE, get_error_details_and_drop(&serialize_result.err));
+  }
+
+  ddog_prof_MaybeError advance_gen_result = ddog_prof_ManagedStringStorage_advance_gen(string_storage.ok);
+
+  if (advance_gen_result.tag == DDOG_PROF_OPTION_ERROR_SOME_ERROR) {
+    raise_error(rb_eRuntimeError, "Failed to advance string storage gen: %"PRIsVALUE, get_error_details_and_drop(&advance_gen_result.some));
+  }
+
+  VALUE encoded_pprof_1 = from_ddog_prof_EncodedProfile(serialize_result.ok);
+
+  result = ddog_prof_Profile_add(
+    &profile.ok,
+    (ddog_prof_Sample) {
+      .locations = (ddog_prof_Slice_Location) { .ptr = locations, .len = 1},
+      .values = (ddog_Slice_I64) {.ptr = metric_values, .len = 8},
+      .labels = (ddog_prof_Slice_Label) { .ptr = labels, .len = 1 }
+    },
+    0
+  );
+
+  if (result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
+    raise_error(rb_eArgError, "Failed to record sample: %"PRIsVALUE, get_error_details_and_drop(&result.err));
+  }
+
+  serialize_result = ddog_prof_Profile_serialize(&profile.ok, &start_timestamp, &finish_timestamp);
+
+  if (serialize_result.tag == DDOG_PROF_PROFILE_SERIALIZE_RESULT_ERR) {
+    raise_error(rb_eArgError, "Failed to serialize: %"PRIsVALUE, get_error_details_and_drop(&serialize_result.err));
+  }
+
+  VALUE encoded_pprof_2 = from_ddog_prof_EncodedProfile(serialize_result.ok);
+
+  ddog_prof_Profile_drop(&profile.ok);
+  ddog_prof_ManagedStringStorage_drop(string_storage.ok);
+
+  return rb_ary_new_from_args(2, encoded_pprof_1, encoded_pprof_2);
+}
+
+static VALUE _native_commit_heap_recordings(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  heap_recorder_commit_recordings_may_lose_gvl(state->heap_recorder);
+
+  return Qtrue;
+}
+
+void recorder_install_on_serialize(VALUE recorder_instance, VALUE thread_context_collector_instance) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  state->thread_context_collector_instance = enforce_thread_context_collector_instance(thread_context_collector_instance);
+}