Serialize register/remove operations.

Author: samuel-williams-shopify

lib/async/container/supervisor/memory_monitor.rb

@@ -30,6 +30,9 @@ def initialize(interval: 10, total_size_limit: nil, memory_sample: false, **opti
 					@options = options
 
 					@processes = Hash.new{|hash, key| hash[key] = Set.new.compare_by_identity}
+					
+					# Queue to serialize cluster modifications to prevent race conditions:
+					@queue = Thread::Queue.new
 				end
 
 				# @attribute [Memory::Leak::Cluster] The cluster of processes being monitored.
@@ -42,6 +45,15 @@ def add(process_id)
 					@cluster.add(process_id, **@options)
 				end
 
+				# Process all pending operations from the queue.
+				#
+				# This ensures cluster modifications are serialized and don't conflict with @cluster.check! iterations.
+				private def process_queue!
+					while operation = @queue.pop(timeout: 0)
+						operation.call
+					end
+				end
+				
 				# Register the connection (worker) with the memory monitor.
 				def register(connection)
 					Console.debug(self, "Registering connection.", connection: connection, state: connection.state)
@@ -50,7 +62,8 @@ def register(connection)
 
 						if connections.empty?
 							Console.debug(self, "Registering process.", child: {process_id: process_id})
-							self.add(process_id)
+							# Queue the cluster modification to avoid race conditions:
+							@queue.push(proc { @cluster.add(process_id, **@options) })
 						end
 
 						connections.add(connection)

test/async/container/memory_monitor.rb

@@ -14,23 +14,283 @@
 	let(:monitor) {subject.new(interval: 1, memory_sample: {duration: 1, timeout: 5})}
 	it_behaves_like Async::Container::Supervisor::AMonitor
 
+	# Mock connection object for testing
+	let(:mock_connection) do
+		Object.new.tap do |connection|
+			def connection.state
+				@state ||= {}
+			end
+		end
+	end
+	
+	with "#register" do
+		it "adds process to cluster when registering first connection" do
+			process_id = 12345
+			mock_connection.state[:process_id] = process_id
+			
+			expect(monitor.cluster).to receive(:add).with(process_id)
+			
+			monitor.register(mock_connection)
+			
+			# Process the queue to execute the queued operation:
+			monitor.send(:process_queue!)
+		end
+		
+		it "does not add process to cluster when connection already exists" do
+			process_id = 12345
+			mock_connection.state[:process_id] = process_id
+			
+			# Register first connection
+			monitor.register(mock_connection)
+			monitor.send(:process_queue!)
+			
+			# Verify process was added
+			expect(monitor.cluster.processes.keys).to be(:include?, process_id)
+			
+			# Register second connection to same process - should not queue another add
+			mock_connection2 = Object.new.tap do |conn|
+				def conn.state
+					@state ||= {process_id: 12345}
+				end
+			end
+			
+			# Count queue size before
+			queue_size_before = monitor.instance_variable_get(:@queue).size
+			
+			monitor.register(mock_connection2)
+			
+			# Queue size should not increase (no new add operation queued)
+			queue_size_after = monitor.instance_variable_get(:@queue).size
+			expect(queue_size_after).to be == queue_size_before
+		end
+		
+		it "handles connection without process_id gracefully" do
+			mock_connection.state.clear
+			
+			# Should not raise an error
+			monitor.register(mock_connection)
+			monitor.send(:process_queue!)
+			
+			# Queue should be empty (no operations queued)
+			expect(monitor.instance_variable_get(:@queue).size).to be == 0
+		end
+	end
+	
+	with "#remove" do
+		it "removes process from cluster when removing last connection" do
+			process_id = 12345
+			mock_connection.state[:process_id] = process_id
+			
+			# Register first
+			monitor.register(mock_connection)
+			monitor.send(:process_queue!)
+			
+			# Remove
+			expect(monitor.cluster).to receive(:remove).with(process_id)
+			
+			monitor.remove(mock_connection)
+			monitor.send(:process_queue!)
+		end
+		
+		it "does not remove process from cluster when other connections exist" do
+			process_id = 12345
+			mock_connection.state[:process_id] = process_id
+			
+			# Register first connection
+			monitor.register(mock_connection)
+			monitor.send(:process_queue!)
+			
+			# Register second connection
+			mock_connection2 = Object.new.tap do |conn|
+				def conn.state
+					@state ||= {process_id: 12345}
+				end
+			end
+			monitor.register(mock_connection2)
+			monitor.send(:process_queue!)
+			
+			# Verify process is still in cluster
+			expect(monitor.cluster.processes.keys).to be(:include?, process_id)
+			
+			# Count queue size before
+			queue_size_before = monitor.instance_variable_get(:@queue).size
+			
+			# Remove first connection - should not queue remove (other connection exists)
+			monitor.remove(mock_connection)
+			
+			# Queue size should not increase
+			queue_size_after = monitor.instance_variable_get(:@queue).size
+			expect(queue_size_after).to be == queue_size_before
+		end
+		
+		it "handles connection without process_id gracefully" do
+			mock_connection.state.clear
+			
+			# Should not raise an error
+			monitor.remove(mock_connection)
+			monitor.send(:process_queue!)
+			
+			# Queue should be empty (no operations queued)
+			expect(monitor.instance_variable_get(:@queue).size).to be == 0
+		end
+	end
+	
+	with "queue serialization" do
+		include Sus::Fixtures::Async::SchedulerContext
+		
+		it "processes queued operations before cluster check" do
+			process_id = 12345
+			mock_connection.state[:process_id] = process_id
+			
+			# Register connection (queues the operation)
+			monitor.register(mock_connection)
+			
+			# Verify process is not yet in cluster
+			expect(monitor.cluster.processes.keys.include?(process_id)).to be == false
+			
+			# Process queue manually to simulate what run does
+			monitor.send(:process_queue!)
+			
+			# Now verify process was added
+			expect(monitor.cluster.processes.keys).to be(:include?, process_id)
+		end
+		
+		it "serializes multiple register/remove operations" do
+			process_id1 = 11111
+			process_id2 = 22222
+			
+			conn1 = Object.new.tap do |c|
+				def c.state
+					@state ||= {process_id: 11111}
+				end
+			end
+			conn2 = Object.new.tap do |c|
+				def c.state
+					@state ||= {process_id: 22222}
+				end
+			end
+			
+			# Queue multiple operations
+			monitor.register(conn1)
+			monitor.register(conn2)
+			monitor.remove(conn1)
+			
+			# Verify both processes are queued but not yet added
+			expect(monitor.cluster.processes.keys.include?(process_id1)).to be == false
+			expect(monitor.cluster.processes.keys.include?(process_id2)).to be == false
+			
+			# Process queue - should execute all operations
+			monitor.send(:process_queue!)
+			
+			# Verify operations were executed: process_id2 should be added, process_id1 should be removed
+			expect(monitor.cluster.processes.keys).to be(:include?, process_id2)
+			expect(monitor.cluster.processes.keys.include?(process_id1)).to be == false
+		end
+	end
+	
+	with "#memory_leak_detected" do
+		let(:monitor_without_sample) {subject.new(interval: 1)}
+		
+		it "kills process when memory leak is detected" do
+			process_id = 12345
+			mock_monitor = Object.new
+			
+			expect(Process).to receive(:kill).with(:INT, process_id)
+			
+			result = monitor_without_sample.memory_leak_detected(process_id, mock_monitor)
+			
+			expect(result).to be == true
+		end
+		
+		it "handles already-dead process gracefully" do
+			process_id = 99999 # Non-existent process
+			mock_monitor = Object.new
+			
+			expect(Process).to receive(:kill).with(:INT, process_id).and_raise(Errno::ESRCH)
+			
+			# Should not raise an error
+			result = monitor_without_sample.memory_leak_detected(process_id, mock_monitor)
+			
+			expect(result).to be == true
+		end
+		
+		it "captures memory sample when enabled" do
+			process_id = 12345
+			mock_monitor = Object.new
+			mock_connection.state[:process_id] = process_id
+			
+			# Register connection
+			monitor.register(mock_connection)
+			monitor.send(:process_queue!)
+			
+			# Mock the connection call
+			expect(mock_connection).to receive(:call).with(
+				do: :memory_sample,
+				duration: 1,
+				timeout: 5
+			).and_return({data: "sample data"})
+			
+			expect(Process).to receive(:kill).with(:INT, process_id)
+			
+			monitor.memory_leak_detected(process_id, mock_monitor)
+		end
+	end
+	
 	with "#run" do
 		include Sus::Fixtures::Async::SchedulerContext
 
 		it "can run the monitor" do
 			task = monitor.run
 			expect(task).to be(:running?)
+		ensure
+			task&.stop
 		end
 
 		it "can handle failures" do
+			# Mock check! to raise an error - the monitor should continue running
 			expect(monitor.cluster).to receive(:check!).and_raise(Errno::ESRCH)
 
 			task = monitor.run
 			expect(task).to be(:running?)
 
-			sleep 1
+			reactor.sleep(0.1)
 
+			# Task should still be running despite errors
 			expect(task).to be(:running?)
+		ensure
+			task&.stop
+		end
+		
+		it "processes queue before each check iteration" do
+			process_id = 12345
+			mock_connection.state[:process_id] = process_id
+			
+			# Register connection (queues operation)
+			monitor.register(mock_connection)
+			
+			# Verify process is not yet in cluster
+			expect(monitor.cluster.processes.keys.include?(process_id)).to be == false
+			
+			# Track check! calls and verify queue processing
+			check_count = 0
+			process_added_before_check = false
+			original_check = monitor.cluster.method(:check!)
+			expect(monitor.cluster).to receive(:check!) do |&block|
+				check_count += 1
+				# On first check, verify the process was added (queue was processed)
+				if check_count == 1
+					process_added_before_check = monitor.cluster.processes.keys.include?(process_id)
+				end
+				original_check.call(&block)
+			end
+			
+			task = monitor.run
+			reactor.sleep(0.1)
+			
+			expect(check_count).to be > 0
+			expect(process_added_before_check).to be_truthy
+		ensure
+			task&.stop
 		end
 	end
 end