Add memory sample support.

Author: samuel-williams-shopify

.rubocop.yml

@@ -34,6 +34,9 @@ Layout/BeginEndAlignment:
   Enabled: true
   EnforcedStyleAlignWith: start_of_line
 
+Layout/RescueEnsureAlignment:
+  Enabled: true
+
 Layout/ElseAlignment:
   Enabled: true
 
@@ -42,6 +45,7 @@ Layout/DefEndAlignment:
 
 Layout/CaseIndentation:
   Enabled: true
+  EnforcedStyle: end
 
 Layout/CommentIndentation:
   Enabled: true

async-container-supervisor.gemspec

@@ -26,5 +26,6 @@ Gem::Specification.new do |spec|
 
 	spec.add_dependency "async-service"
 	spec.add_dependency "io-endpoint"
+	spec.add_dependency "memory", "~> 0.6"
 	spec.add_dependency "memory-leak", "~> 0.5"
 end

bake/async/container/supervisor.rb

@@ -29,6 +29,25 @@ def status
 	end
 end
 
+# Sample memory allocations from a worker over a time period.
+#
+# This is useful for identifying memory leaks by tracking allocations
+# that are retained after garbage collection.
+#
+# @parameter duration [Integer] The duration in seconds to sample for (default: 10).
+# @parameter connection_id [String] The connection ID to target a specific worker.
+def memory_sample(duration: 10, connection_id:)
+	client do |connection|
+		Console.info(self, "Sampling memory from worker...", duration: duration, connection_id: connection_id)
+		
+		# Build the operation request:
+		operation = {do: :memory_sample, duration: duration}
+		
+		# Use the forward operation to proxy the request to a worker:
+		return connection.call(do: :forward, operation: operation, connection_id: connection_id)
+	end
+end
+
 private
 
 def endpoint

context/getting-started.md

@@ -139,13 +139,46 @@ The {ruby Async::Container::Supervisor::MemoryMonitor} will periodically check w
 
 The supervisor can collect various diagnostics from workers on demand:
 
-- **Memory dumps**: Full heap dumps for memory analysis
-- **Thread dumps**: Stack traces of all threads
+- **Memory dumps**: Full heap dumps for memory analysis via `ObjectSpace.dump_all`.
+- **Memory samples**: Lightweight sampling to identify memory leaks.
+- **Thread dumps**: Stack traces of all threads.
 - **Scheduler dumps**: Async fiber hierarchy
 - **Garbage collection profiles**: GC performance data
 
 These can be triggered programmatically or via command-line tools (when available).
 
+#### Memory Leak Diagnosis
+
+To identify memory leaks, you can use the memory sampling feature which is much lighter weight than a full memory dump. It tracks allocations over a time period and focuses on retained objects.
+
+**Using the bake task:**
+
+```bash
+# Sample for 30 seconds and print report to console
+$ bake async:container:supervisor:memory_sample duration=30
+```
+
+**Programmatically:**
+
+```ruby
+# Assuming you have a connection to a worker:
+result = connection.call(do: :memory_sample, duration: 30)
+puts result[:data]
+```
+
+This will sample memory allocations for the specified duration, then force a garbage collection and return a JSON report showing what objects were allocated during that period and retained after GC. Late-lifecycle allocations that are retained are likely memory leaks.
+
+The JSON report includes:
+- `total_allocated`: Total allocated memory and count
+- `total_retained`: Total retained memory and count  
+- `by_gem`: Breakdown by gem/library
+- `by_file`: Breakdown by source file
+- `by_location`: Breakdown by specific file:line locations
+- `by_class`: Breakdown by object class
+- `strings`: String allocation analysis
+
+This is much more efficient than `do: :memory_dump` which uses `ObjectSpace.dump_all` and can be slow and blocking on large heaps. The JSON format also makes it easy to integrate with monitoring and analysis tools.
+
 ## Advanced Usage
 
 ### Custom Monitors

examples/memory-leak/service.rb

@@ -51,7 +51,7 @@ def setup(container)
 			# The interval at which to check for memory leaks.
 			interval: 1,
 			# The total size limit of all processes:
-			maximum_size_limit: 1024 * 1024 * 40, # 40 MB
+			maximum_size_limit: 1024 * 1024 * 1000, # 1000 MB
 		)]
 	end
 end

examples/simple/simple.rb

@@ -46,6 +46,9 @@ def setup(container)
 	include Async::Container::Supervisor::Environment
 
 	monitors do
-		[Async::Container::Supervisor::MemoryMonitor.new(interval: 1, limit: 1024 * 1024 * 400)]
+		[Async::Container::Supervisor::MemoryMonitor.new(
+			interval: 1,
+			limit: 1024 * 1024 * 400
+		)]
 	end
 end

guides/getting-started/readme.md

@@ -139,13 +139,46 @@ The {ruby Async::Container::Supervisor::MemoryMonitor} will periodically check w
 
 The supervisor can collect various diagnostics from workers on demand:
 
-- **Memory dumps**: Full heap dumps for memory analysis
-- **Thread dumps**: Stack traces of all threads
+- **Memory dumps**: Full heap dumps for memory analysis via `ObjectSpace.dump_all`.
+- **Memory samples**: Lightweight sampling to identify memory leaks.
+- **Thread dumps**: Stack traces of all threads.
 - **Scheduler dumps**: Async fiber hierarchy
 - **Garbage collection profiles**: GC performance data
 
 These can be triggered programmatically or via command-line tools (when available).
 
+#### Memory Leak Diagnosis
+
+To identify memory leaks, you can use the memory sampling feature which is much lighter weight than a full memory dump. It tracks allocations over a time period and focuses on retained objects.
+
+**Using the bake task:**
+
+```bash
+# Sample for 30 seconds and print report to console
+$ bake async:container:supervisor:memory_sample duration=30
+```
+
+**Programmatically:**
+
+```ruby
+# Assuming you have a connection to a worker:
+result = connection.call(do: :memory_sample, duration: 30)
+puts result[:data]
+```
+
+This will sample memory allocations for the specified duration, then force a garbage collection and return a JSON report showing what objects were allocated during that period and retained after GC. Late-lifecycle allocations that are retained are likely memory leaks.
+
+The JSON report includes:
+- `total_allocated`: Total allocated memory and count
+- `total_retained`: Total retained memory and count  
+- `by_gem`: Breakdown by gem/library
+- `by_file`: Breakdown by source file
+- `by_location`: Breakdown by specific file:line locations
+- `by_class`: Breakdown by object class
+- `strings`: String allocation analysis
+
+This is much more efficient than `do: :memory_dump` which uses `ObjectSpace.dump_all` and can be slow and blocking on large heaps. The JSON format also makes it easy to integrate with monitoring and analysis tools.
+
 ## Advanced Usage
 
 ### Custom Monitors

lib/async/container/supervisor/connection.rb

@@ -71,6 +71,27 @@ def closed?
 						@queue.closed?
 					end
 
+					# Forward this call to another connection, proxying all responses back.
+					#
+					# This provides true streaming forwarding - intermediate responses flow through
+					# in real-time rather than being buffered.
+					#
+					# @parameter target_connection [Connection] The connection to forward the call to.
+					# @parameter operation [Hash] The operation request to forward (must include :do key).
+					def forward(target_connection, operation)
+						# Forward the operation in an async task to avoid blocking
+						Async do
+							# Make the call to the target connection and stream responses back:
+							Call.call(target_connection, **operation) do |response|
+								# Push each response through our queue:
+								self.push(**response)
+							end
+						ensure
+							# Close our queue to signal completion:
+							@queue.close
+						end
+					end
+					
 					def self.dispatch(connection, target, id, message)
 						Async do
 							call = self.new(connection, id, message)

lib/async/container/supervisor/memory_monitor.rb

@@ -10,15 +10,19 @@ module Async
 	module Container
 		module Supervisor
 			class MemoryMonitor
+				MEMORY_SAMPLE = {duration: 60, timeout: 60+20}
+				
 				# Create a new memory monitor.
 				#
 				# @parameter interval [Integer] The interval at which to check for memory leaks.
 				# @parameter total_size_limit [Integer] The total size limit of all processes, or nil for no limit.
 				# @parameter options [Hash] Options to pass to the cluster when adding processes.
-				def initialize(interval: 10, total_size_limit: nil, **options)
+				def initialize(interval: 10, total_size_limit: nil, memory_sample: MEMORY_SAMPLE, **options)
 					@interval = interval
 					@cluster = Memory::Leak::Cluster.new(total_size_limit: total_size_limit)
 
+					@memory_sample = memory_sample
+					
 					# We use these options when adding processes to the cluster:
 					@options = options
 
@@ -74,6 +78,23 @@ def status(call)
 				# @parameter monitor [Memory::Leak::Monitor] The monitor that detected the memory leak.
 				# @returns [Boolean] True if the process was killed.
 				def memory_leak_detected(process_id, monitor)
+					Console.info(self, "Memory leak detected!", child: {process_id: process_id}, monitor: monitor)
+					
+					if @memory_sample
+						Console.info(self, "Capturing memory sample...", child: {process_id: process_id}, memory_sample: @memory_sample)
+
+						# We are tracking multiple connections to the same process:
+						connections = @processes[process_id]
+						
+						# Try to capture a memory sample:
+						connections.each do |connection|
+							result = connection.call(do: :memory_sample, **@memory_sample)
+							
+							Console.info(self, "Memory sample completed:", child: {process_id: process_id}, result: result)
+						end
+					end
+					
+					# Kill the process gently:
 					Console.info(self, "Killing process!", child: {process_id: process_id})
 					Process.kill(:INT, process_id)
 

lib/async/container/supervisor/server.rb

@@ -3,6 +3,8 @@
 # Released under the MIT License.
 # Copyright, 2025, by Samuel Williams.
 
+require "securerandom"
+
 require_relative "connection"
 require_relative "endpoint"
 require_relative "dispatchable"
@@ -17,15 +19,23 @@ class Server
 				def initialize(monitors: [], endpoint: Supervisor.endpoint)
 					@monitors = monitors
 					@endpoint = endpoint
+					
+					@connections = {}
 				end
 
 				attr :monitors
+				attr :connections
 
 				include Dispatchable
 
 				def do_register(call)
 					call.connection.state.merge!(call.message[:state])
 
+					connection_id = SecureRandom.uuid
+					call.connection.state[:connection_id] = connection_id
+					
+					@connections[connection_id] = call.connection
+					
 					@monitors.each do |monitor|
 						monitor.register(call.connection)
 					rescue => error
@@ -35,6 +45,31 @@ def do_register(call)
 					call.finish
 				end
 
+				# Forward an operation to a worker connection.
+				#
+				# @parameter call [Connection::Call] The call to handle.
+				# @parameter operation [Hash] The operation to forward, must include :do key.
+				# @parameter connection_id [String] The connection ID to target.
+				def do_forward(call)
+					operation = call[:operation]
+					connection_id = call[:connection_id]
+					
+					unless connection_id
+						call.fail(error: "Missing 'connection_id' parameter")
+						return
+					end
+					
+					connection = @connections[connection_id]
+					
+					unless connection
+						call.fail(error: "Connection not found", connection_id: connection_id)
+						return
+					end
+					
+					# Forward the call to the target connection
+					call.forward(connection, operation)
+				end
+				
 				# Restart the current process group, usually including the supervisor and any other processes.
 				#
 				# @parameter signal [Symbol] The signal to send to the process group.
@@ -48,14 +83,26 @@ def do_restart(call)
 				end
 
 				def do_status(call)
+					connections = @connections.map do |connection_id, connection|
+						{
+							connection_id: connection_id,
+							process_id: connection.state[:process_id],
+							state: connection.state,
+						}
+					end
+					
 					@monitors.each do |monitor|
 						monitor.status(call)
 					end
 
-					call.finish
+					call.finish(connections: connections)
 				end
 
 				def remove(connection)
+					if connection_id = connection.state[:connection_id]
+						@connections.delete(connection_id)
+					end
+					
 					@monitors.each do |monitor|
 						monitor.remove(connection)
 					rescue => error