Use for object states.

Author: samuel-williams-shopify

ext/memory/profiler/allocations.c

@@ -12,8 +12,9 @@ static VALUE Memory_Profiler_Allocations = Qnil;
 
 // Helper to mark states table (object => state)
 static int Memory_Profiler_Allocations_states_mark(st_data_t key, st_data_t value, st_data_t arg) {
-	// Don't mark the object key (weak reference - don't keep objects alive)
-	// Mark the state value
+	// We don't want the key to move - we can't rehash the table if it does.
+	rb_gc_mark((VALUE)key);
+	
 	VALUE state = (VALUE)value;
 	rb_gc_mark_movable(state);
 	return ST_CONTINUE;
@@ -27,19 +28,7 @@ static int Memory_Profiler_Allocations_states_foreach(st_data_t key, st_data_t v
 
 // Replace callback for st_foreach_with_replace to update states during compaction
 static int Memory_Profiler_Allocations_states_compact(st_data_t *key, st_data_t *value, st_data_t data, int existing) {
-	// Key is object (VALUE) - we can't update if it moved (would require rehashing/allocation)
-	VALUE old_state = (VALUE)*value;
-	VALUE new_state = rb_gc_location(old_state);
-	
-	// NOTE: We can't update object keys if they moved (would require rehashing/allocation).
-	// This means lookups may fail after compaction for moved objects.
-	// This is acceptable - FREEOBJ will simply not find the state and skip the callback.
-	// The state will be cleaned up when Allocations is freed/cleared.
-	
-	// Update state value if it moved (this is safe, doesn't rehash)
-	if (old_state != new_state) {
-		*value = (st_data_t)new_state;
-	}
+	*value = (st_data_t)rb_gc_location((VALUE)*value);
 
 	return ST_CONTINUE;
 }

ext/memory/profiler/capture.c

@@ -302,10 +302,18 @@ static void Memory_Profiler_Capture_event_callback(VALUE data, void *ptr) {
 		// Skip NEWOBJ if disabled (during callback) to prevent infinite recursion
 		if (capture->paused) return;
 
+		// It's safe to unconditionally call here:
+		object = rb_obj_id(object);
+
 		Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_NEWOBJ, data, klass, object);
 	} else if (event_flag == RUBY_INTERNAL_EVENT_FREEOBJ) {
-		// Always process FREEOBJ to ensure state cleanup
-		Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_FREEOBJ, data, klass, object);
+		// We only care about objects that have been seen before (i.e. have an object ID):
+		if (RB_FL_TEST(object, FL_SEEN_OBJ_ID)) {
+			// It's only safe to call here if the object already has an object ID.
+			object = rb_obj_id(object);
+			
+			Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_FREEOBJ, data, klass, object);
+		}
 	}
 }
 

ext/memory/profiler/events.c

@@ -136,12 +136,7 @@ static void Memory_Profiler_Events_compact_queue(struct Memory_Profiler_Queue *q
 		// Update all VALUEs if they moved during compaction:
 		event->capture = rb_gc_location(event->capture);
 		event->klass = rb_gc_location(event->klass);
-		
-		// For NEWOBJ: update object pointer if it moved
-		// For FREEOBJ: DON'T update (object is being freed, pointer is stale)
-		if (event->type == MEMORY_PROFILER_EVENT_TYPE_NEWOBJ) {
-			event->object = rb_gc_location(event->object);
-		}
+		event->object = rb_gc_location(event->object);
 	}
 }
 

releases.md

@@ -1,5 +1,9 @@
 # Releases
 
+## Unreleased
+
+  - Use `rb_obj_id` for tracking object states to avoid compaction issues.
+
 ## v1.1.11
 
   - Double buffer shared events queues to fix queue corruption.