Fix handling of event queue.

Author: samuel-williams-shopify

ext/memory/profiler/capture.c

@@ -13,6 +13,9 @@
 
 enum {
 	DEBUG = 0,
+
+	// This generates a lot of output:
+	DEBUG_EVENT = 0,
 };
 
 static VALUE Memory_Profiler_Capture = Qnil;
@@ -224,10 +227,7 @@ static void Memory_Profiler_Capture_process_freeobj(VALUE capture_value, VALUE u
 	struct Memory_Profiler_Object_Table_Entry *entry = Memory_Profiler_Object_Table_lookup(capture->states, object);
 
 	if (!entry) {
-		if (DEBUG) fprintf(stderr, "[FREEOBJ] Object not found in table: %p\n", (void*)object);
 		goto done;
-	} else {
-		if (DEBUG) fprintf(stderr, "[FREEOBJ] Object found in table: %p\n", (void*)object);
 	}
 
 	VALUE klass = entry->klass;
@@ -246,7 +246,7 @@ static void Memory_Profiler_Capture_process_freeobj(VALUE capture_value, VALUE u
 	Memory_Profiler_Object_Table_delete_entry(capture->states, entry);
 
 	struct Memory_Profiler_Capture_Allocations *record = Memory_Profiler_Allocations_get(allocations);
-	
+
 	// Increment global free count
 	capture->free_count++;
 
@@ -323,13 +323,12 @@ static void Memory_Profiler_Capture_event_callback(VALUE self, void *ptr) {
 	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
 
 	VALUE object = rb_tracearg_object(trace_arg);
-	
-	// We don't want to track internal non-Object allocations:
-	if (!Memory_Profiler_Capture_trackable_p(object)) return;
-	
 	rb_event_flag_t event_flag = rb_tracearg_event_flag(trace_arg);
 
 	if (event_flag == RUBY_INTERNAL_EVENT_NEWOBJ) {
+		// We don't want to track internal non-Object allocations:
+		if (!Memory_Profiler_Capture_trackable_p(object)) return;
+		
 		// Skip NEWOBJ if disabled (during callback) to prevent infinite recursion
 		if (capture->paused) return;
 
@@ -338,12 +337,10 @@ static void Memory_Profiler_Capture_event_callback(VALUE self, void *ptr) {
 		// Skip if klass is not a Class
 		if (rb_type(klass) != RUBY_T_CLASS) return;
 
-		// Enqueue actual object (not object_id) - queue retains it until processed
-		// Ruby 3.5 compatible: no need for FL_SEEN_OBJ_ID or rb_obj_id
-		if (DEBUG) fprintf(stderr, "[NEWOBJ] Enqueuing event for object: %p\n", (void*)object);
+		if (DEBUG_EVENT) fprintf(stderr, "[NEWOBJ] Enqueuing event for object: %p\n", (void*)object);
 		Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_NEWOBJ, self, klass, object);
 	} else if (event_flag == RUBY_INTERNAL_EVENT_FREEOBJ) {
-		if (DEBUG) fprintf(stderr, "[FREEOBJ] Enqueuing event for object: %p\n", (void*)object);
+		if (DEBUG_EVENT) fprintf(stderr, "[FREEOBJ] Enqueuing event for object: %p\n", (void*)object);
 		Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_FREEOBJ, self, Qnil, object);
 	}
 }
@@ -596,12 +593,19 @@ struct Memory_Profiler_Each_Object_Arguments {
 
 	// The allocations wrapper to filter by (Qnil = no filter).
 	VALUE allocations;
+	
+	// Previous GC state (to restore in ensure handler)
+	int gc_was_enabled;
 };
 
-// Cleanup function to re-enable GC
+// Cleanup function to restore GC state
 static VALUE Memory_Profiler_Capture_each_object_ensure(VALUE arg) {
-	// Re-enable GC (rb_gc_enable returns previous state, but we don't need it)
-	rb_gc_enable();
+	struct Memory_Profiler_Each_Object_Arguments *arguments = (struct Memory_Profiler_Each_Object_Arguments *)arg;
+	
+	// Restore GC state only if it was enabled before
+	if (arguments->gc_was_enabled) {
+		rb_gc_enable();
+	}
 
 	return Qnil;
 }
@@ -619,8 +623,8 @@ static VALUE Memory_Profiler_Capture_each_object_body(VALUE arg) {
 		for (size_t i = 0; i < capture->states->capacity; i++) {
 			struct Memory_Profiler_Object_Table_Entry *entry = &capture->states->entries[i];
 
-			// Skip empty or deleted slots (0 = not set)
-			if (entry->object == 0) {
+			// Skip empty or deleted slots (0 = not set, Qnil = deleted)
+			if (entry->object == 0 || entry->object == Qnil) {
 				continue;
 			}
 
@@ -661,31 +665,31 @@ static VALUE Memory_Profiler_Capture_each_object(int argc, VALUE *argv, VALUE se
 
 	RETURN_ENUMERATOR(self, argc, argv);
 
-	// Disable GC to prevent objects from being collected during iteration
-	rb_gc_disable();
-	
-	// Process all pending events to clean up stale entries
-	// At this point, all remaining objects in the table should be valid
-	Memory_Profiler_Events_process_all();
-	
 	// If class provided, look up its allocations wrapper
 	VALUE allocations = Qnil;
 	if (!NIL_P(klass)) {
 		st_data_t allocations_data;
 		if (st_lookup(capture->tracked, (st_data_t)klass, &allocations_data)) {
 			allocations = (VALUE)allocations_data;
 		} else {
-			// Class not tracked - nothing to iterate
-			// Re-enable GC before returning
-			rb_gc_enable();
+			// Class not tracked - nothing to iterate:
 			return self;
 		}
 	}
+
+	// Disable GC to prevent objects from being collected during iteration
+	// rb_gc_disable returns the previous state (non-zero = was enabled, 0 = was disabled)
+	int gc_was_enabled = RTEST(rb_gc_disable());
+	
+	// Process all pending events to clean up stale entries:
+	Memory_Profiler_Events_process_all();
+	// At this point, all remaining objects in the table should be valid.
 
 	// Setup arguments for iteration
 	struct Memory_Profiler_Each_Object_Arguments arguments = {
 		.self = self,
-		.allocations = allocations
+		.allocations = allocations,
+		.gc_was_enabled = gc_was_enabled
 	};
 
 	// Use rb_ensure to guarantee cleanup even if exception is raised

ext/memory/profiler/events.c

@@ -19,6 +19,9 @@ struct Memory_Profiler_Events {
 	// Double-buffered event queues (contains events from all Capture instances).
 	struct Memory_Profiler_Queue queues[2];
 	struct Memory_Profiler_Queue *available, *processing;
+	
+	// Guard flag to prevent recursive processing (0 = not processing, 1 = processing)
+	int processing_flag;
 
 	// Postponed job handle for processing the queue.
 	// Postponed job handles are an extremely limited resource, so we only register one global event queue.
@@ -59,6 +62,7 @@ static VALUE Memory_Profiler_Events_new(void) {
 	// Start with queues[0] available for incoming events, queues[1] for processing (initially empty):
 	events->available = &events->queues[0];
 	events->processing = &events->queues[1];
+	events->processing_flag = 0;
 
 	// Pre-register the single postponed job for processing the queue:
 	events->postponed_job_handle = rb_postponed_job_preregister(0,
@@ -86,8 +90,6 @@ struct Memory_Profiler_Events* Memory_Profiler_Events_instance(void) {
 		// Pin the global events object so it's never GC'd:
 		rb_gc_register_mark_object(instance);
 
-		if (DEBUG) fprintf(stderr, "Global event queue system initialized and pinned\n");
-		
 		TypedData_Get_Struct(instance, struct Memory_Profiler_Events, &Memory_Profiler_Events_type, events);
 	}
 
@@ -195,14 +197,16 @@ int Memory_Profiler_Events_enqueue(
 		RB_OBJ_WRITE(events->self, &event->klass, klass);
 		RB_OBJ_WRITE(events->self, &event->object, object);
 
-		if (DEBUG) fprintf(stderr, "Queued %s to available queue, size: %zu\n", 
-			Memory_Profiler_Event_Type_name(type), events->available->count);
+		if (DEBUG) {
+			fprintf(stderr, "[EVENTS] Enqueued %s: object=%p available_count=%zu processing_flag=%d\n", 
+				Memory_Profiler_Event_Type_name(type), (void*)object, events->available->count, events->processing_flag);
+		}
 
 		rb_postponed_job_trigger(events->postponed_job_handle);
 		// Success:
 		return 1;
 	}
-
+	
 	// Queue full:
 	return 0;
 }
@@ -211,6 +215,12 @@ int Memory_Profiler_Events_enqueue(
 // Public API function - called from Capture stop() to ensure all events are processed.
 void Memory_Profiler_Events_process_all(void) {
 	struct Memory_Profiler_Events *events = Memory_Profiler_Events_instance();
+	
+	// Explicitly prevent re-entrancy here:
+	if (events->processing_flag) {
+		rb_raise(rb_eRuntimeError, "Recursive call detected!");
+	}
+	
 	Memory_Profiler_Events_process_queue((void *)events);
 }
 
@@ -228,17 +238,39 @@ static VALUE Memory_Profiler_Events_process_event_protected(VALUE arg) {
 static void Memory_Profiler_Events_process_queue(void *arg) {
 	struct Memory_Profiler_Events *events = (struct Memory_Profiler_Events *)arg;
 
+	// Check for recursive call - this would break double buffering!
+	if (events->processing_flag) {
+		// Explicitly allow re-entrancy here, as the postponed job could be triggered during `process_all`.
+		return;
+	}
+	
+	// Set processing flag to prevent recursion
+	events->processing_flag = 1;
+	
+	if (DEBUG) {
+		fprintf(stderr, "[EVENTS] process_queue START: available_count=%zu processing_count=%zu\n",
+			events->available->count, events->processing->count);
+	}
+	
 	// Swap the queues: available becomes processing, and the old processing queue (now empty) becomes available. This allows new events to continue enqueueing to the new available queue while we process.
 	struct Memory_Profiler_Queue *queue_to_process = events->available;
 	events->available = events->processing;
 	events->processing = queue_to_process;
 
-	if (DEBUG) fprintf(stderr, "Processing event queue: %zu events\n", events->processing->count);
+	if (DEBUG) {
+		fprintf(stderr, "[EVENTS] Queues swapped: processing_count=%zu (was available), available_count=%zu (was processing)\n",
+			events->processing->count, events->available->count);
+	}
 
 	// Process all events in order (maintains NEWOBJ before FREEOBJ for same object):
 	for (size_t i = 0; i < events->processing->count; i++) {
 		struct Memory_Profiler_Event *event = Memory_Profiler_Queue_at(events->processing, i);
 
+		if (DEBUG) {
+			fprintf(stderr, "[EVENTS] Processing event[%zu]: type=%s object=%p capture=%p\n",
+				i, Memory_Profiler_Event_Type_name(event->type), (void*)event->object, (void*)event->capture);
+		}
+		
 		// Process event with rb_protect to catch any exceptions:
 		int state = 0;
 		rb_protect(Memory_Profiler_Events_process_event_protected, (VALUE)event, &state);
@@ -249,13 +281,28 @@ static void Memory_Profiler_Events_process_queue(void *arg) {
 			rb_set_errinfo(Qnil);
 		}
 
+		if (DEBUG) {
+			fprintf(stderr, "[EVENTS] Processed event[%zu]: type=%s object=%p (exception=%d)\n",
+				i, Memory_Profiler_Event_Type_name(event->type), (void*)event->object, state);
+		}
+		
 		// Clear this event after processing to prevent marking stale data if GC runs:
 		event->type = MEMORY_PROFILER_EVENT_TYPE_NONE;
 		RB_OBJ_WRITE(events->self, &event->capture, Qnil);
 		RB_OBJ_WRITE(events->self, &event->klass, Qnil);
 		RB_OBJ_WRITE(events->self, &event->object, Qnil);
 	}
 
+	// Save count before clearing for logging
+	size_t processed_count = events->processing->count;
+	
 	// Clear the processing queue (which is now empty logically):
 	Memory_Profiler_Queue_clear(events->processing);
+	
+	// Clear processing flag
+	events->processing_flag = 0;
+	
+	if (DEBUG) {
+		fprintf(stderr, "[EVENTS] process_queue END: processed %zu events\n", processed_count);
+	}
 }

ext/memory/profiler/table.c

@@ -365,7 +365,7 @@ void Memory_Profiler_Object_Table_compact(struct Memory_Profiler_Object_Table *t
 	table->count = 0;
 	table->tombstones = 0;  // Compaction clears tombstones
 
-	// Reinsert all entries with new hash values
+	// Reinsert all entries with new hash values:
 	for (size_t i = 0; i < temp_count; i++) {
 		int found;
 		size_t index = find_entry(table->entries, table->capacity, temp_entries[i].object, &found, NULL, "compact");

releases.md

@@ -1,5 +1,10 @@
 # Releases
 
+## Unreleased
+
+  - Ensure all `FREEOBJ` events are enqueued.
+  - Allow limited re-entrancy in queue processing.
+
 ## v1.6.1
 
   - Add `track_all` option to `Sampler`.