Explicit track all allocations.

Author: samuel-williams-shopify

examples/churn.rb

@@ -121,7 +121,9 @@
 puts
 
 puts "Verification:"
-expected = RETAINED_COUNT
+# Expected count accounts for: Hash + String (hash value) + String + Array = 4 objects per iteration
+# So: (RETAINED_COUNT / 3) * 4 = RETAINED_COUNT * 4 / 3
+expected = (RETAINED_COUNT * 4 / 3).round
 tolerance = expected * 0.1  # Allow 10% variance due to GC timing
 diff = (total_live - expected).abs
 

ext/memory/profiler/capture.c

@@ -28,6 +28,9 @@ struct Memory_Profiler_Capture {
 	// Should we queue callbacks? (temporarily disabled during queue processing).
 	int paused;
 
+	// Should we automatically track all classes? (if false, only explicitly tracked classes are tracked).
+	int track_all;
+
 	// Tracked classes: class => VALUE (wrapped Memory_Profiler_Capture_Allocations).
 	st_table *tracked;
 
@@ -163,21 +166,18 @@ static void Memory_Profiler_Capture_process_newobj(VALUE self, VALUE klass, VALU
 	// Pause the capture to prevent infinite loop:
 	capture->paused += 1;
 
-	// Increment global new count:
-	capture->new_count++;
-	
-	// Look up or create allocations record for this class:
+	// Look up allocations record for this class:
 	st_data_t allocations_data;
 	VALUE allocations;
 	struct Memory_Profiler_Capture_Allocations *record;
 
 	if (st_lookup(capture->tracked, (st_data_t)klass, &allocations_data)) {
-		// Existing record
+		// Existing record - class is explicitly tracked
 		allocations = (VALUE)allocations_data;
 		record = Memory_Profiler_Allocations_get(allocations);
 		record->new_count++;
-	} else {
-		// First time seeing this class, create record automatically
+	} else if (capture->track_all) {
+		// First time seeing this class, create record automatically (if track_all is enabled)
 		record = ALLOC(struct Memory_Profiler_Capture_Allocations);
 		record->callback = Qnil;
 		record->new_count = 1;
@@ -187,8 +187,15 @@ static void Memory_Profiler_Capture_process_newobj(VALUE self, VALUE klass, VALU
 		st_insert(capture->tracked, (st_data_t)klass, (st_data_t)allocations);
 		RB_OBJ_WRITTEN(self, Qnil, klass);
 		RB_OBJ_WRITTEN(self, Qnil, allocations);
+	} else {
+		// track_all disabled and class not explicitly tracked - skip this allocation entirely
+		capture->paused -= 1;
+		return;
 	}
 
+	// Increment global new count (only if we're tracking this class):
+	capture->new_count++;
+	
 	VALUE data = Qnil;
 	if (!NIL_P(record->callback)) {
 		data = rb_funcall(record->callback, rb_intern("call"), 3, klass, sym_newobj, Qnil);
@@ -367,9 +374,10 @@ static VALUE Memory_Profiler_Capture_alloc(VALUE klass) {
 	capture->new_count = 0;
 	capture->free_count = 0;
 
-	// Initialize state flags - not running, callbacks disabled
+	// Initialize state flags - not running, callbacks disabled, track_all disabled by default
 	capture->running = 0;
 	capture->paused = 0;
+	capture->track_all = 0;
 
 	// Global event queue system will auto-initialize on first use (lazy initialization)
 
@@ -381,6 +389,24 @@ static VALUE Memory_Profiler_Capture_initialize(VALUE self) {
 	return self;
 }
 
+// Get track_all setting
+static VALUE Memory_Profiler_Capture_track_all_get(VALUE self) {
+	struct Memory_Profiler_Capture *capture;
+	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
+	
+	return capture->track_all ? Qtrue : Qfalse;
+}
+
+// Set track_all setting
+static VALUE Memory_Profiler_Capture_track_all_set(VALUE self, VALUE value) {
+	struct Memory_Profiler_Capture *capture;
+	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
+	
+	capture->track_all = RTEST(value) ? 1 : 0;
+	
+	return value;
+}
+
 // Start capturing allocations
 static VALUE Memory_Profiler_Capture_start(VALUE self) {
 	struct Memory_Profiler_Capture *capture;
@@ -744,6 +770,8 @@ void Init_Memory_Profiler_Capture(VALUE Memory_Profiler)
 	rb_define_alloc_func(Memory_Profiler_Capture, Memory_Profiler_Capture_alloc);
 
 	rb_define_method(Memory_Profiler_Capture, "initialize", Memory_Profiler_Capture_initialize, 0);
+	rb_define_method(Memory_Profiler_Capture, "track_all", Memory_Profiler_Capture_track_all_get, 0);
+	rb_define_method(Memory_Profiler_Capture, "track_all=", Memory_Profiler_Capture_track_all_set, 1);
 	rb_define_method(Memory_Profiler_Capture, "start", Memory_Profiler_Capture_start, 0);
 	rb_define_method(Memory_Profiler_Capture, "stop", Memory_Profiler_Capture_stop, 0);
 	rb_define_method(Memory_Profiler_Capture, "track", Memory_Profiler_Capture_track, -1);  // -1 to accept block

lib/memory/profiler/sampler.rb

@@ -104,6 +104,7 @@ def initialize(depth: 4, filter: nil, increases_threshold: 10, prune_limit: 5, p
 				@gc = gc
 
 				@capture = Capture.new
+				@capture.track_all = true
 				@call_trees = {}
 				@samples = {}
 			end

test/memory/profiler/allocations.rb

@@ -50,6 +50,7 @@
 	with "integration with Capture" do
 		it "returns correct JSON from captured allocations" do
 			capture = Memory::Profiler::Capture.new
+			capture.track_all = true
 			capture.start
 
 			# Create some objects

test/memory/profiler/capture.rb

@@ -55,6 +55,7 @@
 
 	with "#retained_count_of" do
 		it "tracks allocations for a class" do
+			capture.track_all = true
 			capture.start
 
 			initial_count = capture.retained_count_of(Hash)
@@ -226,6 +227,9 @@
 			capture1 = Memory::Profiler::Capture.new
 			capture2 = Memory::Profiler::Capture.new
 
+			capture1.track_all = true
+			capture2.track_all = true
+			
 			capture1.track(Hash)
 			capture2.track(Array)
 
@@ -350,6 +354,7 @@
 
 	with "#new_count, #free_count, #retained_count" do
 		it "tracks total allocations across all classes" do
+			capture.track_all = true
 			capture.start
 
 			# Allocate various objects
@@ -372,6 +377,7 @@
 		it "tracks frees when objects are collected" do
 			GC.start
 
+			capture.track_all = true
 			capture.start
 
 			initial_new = capture.new_count
@@ -395,6 +401,7 @@
 		it "retained_count reflects objects still alive" do
 			GC.start
 
+			capture.track_all = true
 			capture.start
 
 			# Allocate and retain some objects
@@ -418,6 +425,7 @@
 		end
 
 		it "counts reset after clear" do
+			capture.track_all = true
 			capture.start
 
 			20.times{Hash.new}