Fix FREEOBJ queue handling using postponed job.

Author: samuel-williams-shopify

examples/basic_tracking.rb

@@ -48,7 +48,7 @@
 puts "-" * 60
 
 sampler = Memory::Profiler::Sampler.new(depth: 10)
-sampler.track_with_analysis(Hash)
+sampler.track(Hash)
 sampler.start
 
 puts "Tracking Hash with call path analysis enabled"

ext/memory/profiler/allocations.c

@@ -12,8 +12,9 @@ static VALUE Memory_Profiler_Allocations = Qnil;
 
 // Helper to mark object_states table values
 static int Memory_Profiler_Allocations_object_states_mark(st_data_t key, st_data_t value, st_data_t arg) {
-	VALUE object = (VALUE)key;
-	rb_gc_mark_movable(object);
+	// Don't mark the object, we use it as a key but we never use it as an object, and we don't want to retain it.
+	// VALUE object = (VALUE)key;
+	// rb_gc_mark_movable(object);
 
 	VALUE state = (VALUE)value;
 	if (!NIL_P(state)) {

ext/memory/profiler/capture.c

@@ -14,6 +14,7 @@
 enum {
 	DEBUG = 0,
 	DEBUG_FREED_QUEUE = 0,
+	DEBUG_STATE = 0,
 };
 
 static VALUE Memory_Profiler_Capture = Qnil;
@@ -198,9 +199,7 @@ static void Memory_Profiler_Capture_process_freed_queue(void *arg) {
 	struct Memory_Profiler_Capture *capture;
 	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
 
-	if (DEBUG_FREED_QUEUE) {
-		fprintf(stderr, "Processing freed queue with %zu entries\n", capture->freed_queue.count);
-	}
+	if (DEBUG_FREED_QUEUE) fprintf(stderr, "Processing freed queue with %zu entries\n", capture->freed_queue.count);
 
 	// Process all freed objects in the queue
 	for (size_t i = 0; i < capture->freed_queue.count; i++) {
@@ -245,6 +244,9 @@ static void Memory_Profiler_Capture_newobj_handler(VALUE self, struct Memory_Pro
 				if (!record->object_states) {
 					record->object_states = st_init_numtable();
 				}
+				
+				if (DEBUG_STATE) fprintf(stderr, "Memory_Profiler_Capture_newobj_handler: Inserting state for object: %p (%s)\n", (void *)object, rb_class2name(klass));
+
 				st_insert(record->object_states, (st_data_t)object, (st_data_t)state);
 				// Notify GC about the state VALUE stored in the table
 				RB_OBJ_WRITTEN(self, Qnil, state);
@@ -281,22 +283,29 @@ static void Memory_Profiler_Capture_freeobj_handler(VALUE self, struct Memory_Pr
 
 		// If we have a callback and detailed tracking, queue the freeobj for later processing
 		if (!NIL_P(record->callback) && record->object_states) {
+			if (DEBUG_STATE) fprintf(stderr, "Memory_Profiler_Capture_freeobj_handler: Looking up state for object: %p\n", (void *)object);
+
 			// Look up state stored during NEWOBJ
 			st_data_t state_data;
 			if (st_delete(record->object_states, (st_data_t *)&object, &state_data)) {
+				if (DEBUG_STATE) fprintf(stderr, "Found state for object: %p\n", (void *)object);
 				VALUE state = (VALUE)state_data;
 
 				// Push a new item onto the queue (returns pointer to write to)
 				// NOTE: realloc is safe during GC (doesn't trigger Ruby allocation)
 				struct Memory_Profiler_Queue_Item *freed = Memory_Profiler_Queue_push(&capture->freed_queue);
 				if (freed) {
+					if (DEBUG_FREED_QUEUE) fprintf(stderr, "Queued freed object, queue size now: %zu/%zu\n", capture->freed_queue.count, capture->freed_queue.capacity);
 					// Write directly to the allocated space
 					freed->klass = klass;
 					freed->allocations = allocations;
 					freed->state = state;
 
 					// Trigger postponed job to process the queue after GC
+					if (DEBUG_FREED_QUEUE) fprintf(stderr, "Triggering postponed job to process the queue after GC\n");
 					rb_postponed_job_trigger(capture->postponed_job_handle);
+				} else {
+					if (DEBUG_FREED_QUEUE) fprintf(stderr, "Failed to queue freed object, out of memory\n");
 				}
 				// If push failed (out of memory), silently drop this freeobj event
 			}
@@ -455,6 +464,7 @@ static VALUE Memory_Profiler_Capture_stop(VALUE self) {
 // Add a class to track with optional callback
 // Usage: track(klass) or track(klass) { |obj, klass| ... }
 // Callback can call caller_locations with desired depth
+// Returns the Allocations object for the tracked class
 static VALUE Memory_Profiler_Capture_track(int argc, VALUE *argv, VALUE self) {
 	struct Memory_Profiler_Capture *capture;
 	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
@@ -463,8 +473,10 @@ static VALUE Memory_Profiler_Capture_track(int argc, VALUE *argv, VALUE self) {
 	rb_scan_args(argc, argv, "1&", &klass, &callback);
 
 	st_data_t allocations_data;
+	VALUE allocations;
+	
 	if (st_lookup(capture->tracked_classes, (st_data_t)klass, &allocations_data)) {
-		VALUE allocations = (VALUE)allocations_data;
+		allocations = (VALUE)allocations_data;
 		struct Memory_Profiler_Capture_Allocations *record = Memory_Profiler_Allocations_get(allocations);
 		RB_OBJ_WRITE(self, &record->callback, callback);
 	} else {
@@ -475,7 +487,7 @@ static VALUE Memory_Profiler_Capture_track(int argc, VALUE *argv, VALUE self) {
 		record->object_states = NULL;
 
 		// Wrap the record in a VALUE
-		VALUE allocations = Memory_Profiler_Allocations_wrap(record);
+		allocations = Memory_Profiler_Allocations_wrap(record);
 
 		st_insert(capture->tracked_classes, (st_data_t)klass, (st_data_t)allocations);
 		// Notify GC about the class VALUE stored as key in the table
@@ -488,7 +500,7 @@ static VALUE Memory_Profiler_Capture_track(int argc, VALUE *argv, VALUE self) {
 		}
 	}
 
-	return self;
+	return allocations;
 }
 
 // Stop tracking a class
@@ -578,6 +590,19 @@ static VALUE Memory_Profiler_Capture_each(VALUE self) {
 	return self;
 }
 
+// Get allocations for a specific class
+static VALUE Memory_Profiler_Capture_aref(VALUE self, VALUE klass) {
+	struct Memory_Profiler_Capture *capture;
+	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
+	
+	st_data_t allocations_data;
+	if (st_lookup(capture->tracked_classes, (st_data_t)klass, &allocations_data)) {
+		return (VALUE)allocations_data;
+	}
+	
+	return Qnil;
+}
+
 void Init_Memory_Profiler_Capture(VALUE Memory_Profiler)
 {
 	// Initialize event symbols
@@ -597,6 +622,7 @@ void Init_Memory_Profiler_Capture(VALUE Memory_Profiler)
 	rb_define_method(Memory_Profiler_Capture, "tracking?", Memory_Profiler_Capture_tracking_p, 1);
 	rb_define_method(Memory_Profiler_Capture, "count_for", Memory_Profiler_Capture_count_for, 1);
 	rb_define_method(Memory_Profiler_Capture, "each", Memory_Profiler_Capture_each, 0);
+	rb_define_method(Memory_Profiler_Capture, "[]", Memory_Profiler_Capture_aref, 1);
 	rb_define_method(Memory_Profiler_Capture, "clear", Memory_Profiler_Capture_clear, 0);
 
 	// Initialize Allocations class

lib/memory/profiler/sampler.rb

@@ -141,7 +141,7 @@ def sample!
 					if sample.sample!(count)
 						# Check if we should enable detailed tracking
 						if sample.increases >= @increases_threshold && !@call_trees.key?(klass)
-							track_with_analysis_internal(klass, allocations)
+							track(klass, allocations)
 						end
 
 						# Notify about growth if block given
@@ -150,13 +150,19 @@ def sample!
 				end
 			end
 
-			# Internal: Enable tracking with analysis using allocations object
-			private def track_with_analysis_internal(klass, allocations)
+			# Start tracking with call path analysis.
+			#
+			# @parameter klass [Class] The class to track with detailed analysis.
+			def track(klass, allocations = nil)
+				# Track the class and get the allocations object
+				allocations ||= @capture.track(klass)
+				
+				# Set up call tree for this class
 				tree = @call_trees[klass] = CallTree.new
 				depth = @depth
 				filter = @filter
 
-				# Register callback on allocations object with new signature:
+				# Register callback on allocations object:
 				# - On :newobj - returns state (leaf node) which C extension stores
 				# - On :freeobj - receives state back from C extension
 				allocations.track do |klass, event, state|
@@ -166,43 +172,16 @@ def sample!
 						locations = caller_locations(1, depth)
 						filtered = locations.select(&filter)
 						unless filtered.empty?
-							# Record returns the leaf node - return it so C can store it
+							# Record returns the leaf node - return it so C can store it:
 							tree.record(filtered)
 						end
-						# Return nil or the node - C will store whatever we return
+						# Return nil or the node - C will store whatever we return.
 					when :freeobj
-						# Decrement using the state (leaf node) passed back from C
-						if state
-							state.decrement_path!
-						end
+						# Decrement using the state (leaf node) passed back from then native extension:
+						state&.decrement_path!
 					end
 				rescue Exception => error
-					warn "Error in track_with_analysis_internal: #{error.message}\n#{error.backtrace.join("\n")}"
-				end
-			end
-			
-			# Start tracking allocations for a class (count only).
-			def track(klass)
-				return if @capture.tracking?(klass)
-				
-				@capture.track(klass)
-			end
-			
-			# Start tracking with call path analysis.
-			#
-			# @parameter klass [Class] The class to track with detailed analysis.
-			def track_with_analysis(klass)
-				# Track the class if not already tracked
-				unless @capture.tracking?(klass)
-					@capture.track(klass)
-				end
-				
-				# Enable analysis by setting callback on the allocations object
-				@capture.each do |tracked_klass, allocations|
-					if tracked_klass == klass
-						track_with_analysis_internal(klass, allocations)
-						break
-					end
+					warn "Error in allocation tracking: #{error.message}\n#{error.backtrace.join("\n")}"
 				end
 			end
 
@@ -280,11 +259,7 @@ def stop!
 			private
 
 			def default_filter
-				->(location) {path = location.path
-																		!path.include?("/gems/") && 
-																							!path.include?("/ruby/") &&
-																							!path.start_with?("(eval)")
-				}
+				->(location) {!location.path.match?(%r{/(gems|ruby)/|\A\(eval\)})}
 			end
 		end
 	end