Simplify container supervision using Async

Assisted-By: devx/3236e566-7538-432e-a30a-2bdf37265ed4

Author: samuel-williams-shopify

fixtures/async/container/a_container.rb

@@ -73,6 +73,38 @@ module Container
 				end
 			end
 
+			with "Async compatibility" do
+				it "can wait inside an Async task after spawning outside Async" do
+					input, output = IO.pipe
+					
+					container.spawn do
+						output.write(".")
+					end
+					
+					Async do
+						container.wait
+					end
+					
+					output.close
+					expect(input.read).to be == "."
+				end
+				
+				it "can spawn and wait inside the same Async task" do
+					input, output = IO.pipe
+					
+					Async do
+						container.spawn do
+							output.write(".")
+						end
+						
+						container.wait
+					end
+					
+					output.close
+					expect(input.read).to be == "."
+				end
+			end
+			
 			it "should be blocking" do
 				skip "Fiber.blocking? is not supported!" unless Fiber.respond_to?(:blocking?)
 

lib/async/container/channel.rb

@@ -25,12 +25,14 @@ def initialize(timeout: 1.0)
 
 			# Close the input end of the pipe.
 			def close_read
-				@in.close
+				@in.close unless @in.closed?
+			rescue IOError, Errno::EBADF
 			end
 
 			# Close the output end of the pipe.
 			def close_write
-				@out.close
+				@out.close unless @out.closed?
+			rescue IOError, Errno::EBADF
 			end
 
 			# Close both ends of the pipe.

lib/async/container/forked.rb

@@ -101,23 +101,25 @@ def exec(*arguments, ready: true, **options)
 				def self.fork(**options)
 					# $stderr.puts fork: caller
 					self.new(**options) do |process|
-						::Process.fork do
-							# We use `Thread.current.raise(...)` so that exceptions are filtered through `Thread.handle_interrupt` correctly.
-							Signal.trap(:INT){::Thread.current.raise(Interrupt)}
-							Signal.trap(:TERM){::Thread.current.raise(Interrupt)}  # Same as SIGINT.
-							Signal.trap(:HUP){::Thread.current.raise(Restart)}
-							
-							# This could be a configuration option:
-							::Thread.handle_interrupt(SignalException => :immediate) do
-								yield Instance.for(process)
-							rescue Interrupt
-								# Graceful exit.
-							rescue Exception => error
-								Console.error(self, error)
+						::Thread.new do
+							::Process.fork do
+								# We use `Thread.current.raise(...)` so that exceptions are filtered through `Thread.handle_interrupt` correctly.
+								Signal.trap(:INT){::Thread.current.raise(Interrupt)}
+								Signal.trap(:TERM){::Thread.current.raise(Interrupt)}  # Same as SIGINT.
+								Signal.trap(:HUP){::Thread.current.raise(Restart)}
 
-								exit!(1)
+								# This could be a configuration option:
+								::Thread.handle_interrupt(SignalException => :immediate) do
+									yield Instance.for(process)
+								rescue Interrupt
+									# Graceful exit.
+								rescue Exception => error
+									Console.error(self, error)
+									
+									exit!(1)
+								end
 							end
-						end
+						end.value
 					end
 				end
 
@@ -151,6 +153,21 @@ def initialize(name: nil, **options)
 					self.close_write
 				end
 
+				# A minimal status for children reaped outside this object.
+				class Status
+					def success?
+						true
+					end
+					
+					def to_i
+						0
+					end
+					
+					def to_s
+						"\#<#{self.class} success>"
+					end
+				end
+				
 				# Convert the child process to a hash, suitable for serialization.
 				#
 				# @returns [Hash] The request as a hash.
@@ -236,24 +253,28 @@ def restart!
 				# @parameter timeout [Numeric | Nil] Maximum time to wait before forceful termination.
 				# @returns [::Process::Status] The process exit status.
 				def wait(timeout = 0.1)
-					if @pid && @status.nil?
-						Console.debug(self, "Waiting for process to exit...", child: {process_id: @pid}, timeout: timeout)
-						
-						_, @status = ::Process.wait2(@pid, ::Process::WNOHANG)
-						
-						if @status.nil?
-							sleep(timeout) if timeout
+					begin
+						if @pid && @status.nil?
+							Console.debug(self, "Waiting for process to exit...", child: {process_id: @pid}, timeout: timeout)
 
 							_, @status = ::Process.wait2(@pid, ::Process::WNOHANG)
 
 							if @status.nil?
-								Console.warn(self, "Process is blocking, sending kill signal...", child: {process_id: @pid}, timeout: timeout)
-								self.kill!
+								sleep(timeout) if timeout
+								
+								_, @status = ::Process.wait2(@pid, ::Process::WNOHANG)
 
-								# Wait for the process to exit:
-								_, @status = ::Process.wait2(@pid)
+								if @status.nil?
+									Console.warn(self, "Process is blocking, sending kill signal...", child: {process_id: @pid}, timeout: timeout)
+									self.kill!
+									
+									# Wait for the process to exit:
+									_, @status = ::Process.wait2(@pid)
+								end
 							end
 						end
+					rescue Errno::ECHILD
+						@status = Status.new
 					end
 
 					Console.debug(self, "Process exited.", child: {process_id: @pid, status: @status})

lib/async/container/generic.rb

@@ -5,6 +5,7 @@
 # Copyright, 2025, by Marc-André Cournoyer.
 
 require "etc"
+require "io/wait"
 require "async/clock"
 
 require_relative "group"
@@ -144,8 +145,6 @@ def wait_until_ready
 						end
 					end
 
-					self.sleep
-					
 					if self.status?(:ready)
 						Console.debug(self) do |buffer|
 							buffer.puts "All ready:"
@@ -156,6 +155,8 @@ def wait_until_ready
 
 						return true
 					end
+					
+					self.sleep
 				end
 			end
 
@@ -222,13 +223,22 @@ def spawn(name: nil, restart: false, key: nil, health_check_timeout: nil, startu
 				end
 
 				@statistics.spawn!
+				started = ::Thread::Queue.new
 
-				fiber do
+				@group.supervise do
+					first = true
+					
 					until @stopping
 						Console.debug(self, "Starting child...", child: {key: key, name: name, restart: restart, health_check_timeout: health_check_timeout}, statistics: @statistics)
 
 						child = self.start(name, &block)
 						state = insert(key, child)
+						@group.insert(child)
+						
+						if first
+							started << true
+							first = false
+						end
 
 						# Notify policy of spawn
 						begin
@@ -247,36 +257,11 @@ def spawn(name: nil, restart: false, key: nil, health_check_timeout: nil, startu
 						status = nil
 
 						begin
-							status = @group.wait_for(child) do |message|
-								case message
-								when :health_check!
-									if state[:ready]
-										# If a health check timeout is specified, we will monitor the child process and terminate it if it does not update its state within the specified time.
-										if health_check_timeout
-											if health_check_timeout < age_clock.total
-												health_check_failed(child, age_clock, health_check_timeout)
-											end
-										end
-									else
-										# If a startup timeout is specified, we will monitor the child process and terminate it if it does not become ready within the specified time.
-										if startup_timeout
-											if startup_timeout < age_clock.total
-												startup_failed(child, age_clock, startup_timeout)
-											end
-										end
-									end
-								else
-									state.update(message)
-									
-									# Reset the age clock if the child has become ready:
-									if state[:ready]
-										age_clock&.reset!
-									end
-								end
-							end
+							status = wait_for(child, state, age_clock, health_check_timeout, startup_timeout)
 						rescue => error
 							Console.error(self, "Error during child process management!", exception: error, stopping: @stopping)
 						ensure
+							@group.delete(child)
 							delete(key, child)
 						end
 
@@ -300,7 +285,17 @@ def spawn(name: nil, restart: false, key: nil, health_check_timeout: nil, startu
 							break
 						end
 					end
-				end.resume
+				rescue => error
+					started << error if first
+					raise
+				ensure
+					started << false if first
+				end
+				
+				case result = started.pop
+				when Exception
+					raise result
+				end
 
 				return true
 			end
@@ -363,6 +358,78 @@ def key?(key)
 
 			protected
 
+			def wait_for(child, state, age_clock, health_check_timeout, startup_timeout)
+				while true
+					timeout = next_timeout(state, age_clock, health_check_timeout, startup_timeout)
+					
+					if wait_readable(child, timeout)
+						if message = child.receive
+							state.update(message)
+							
+							# Reset the age clock if the child has become ready:
+							if state[:ready]
+								age_clock&.reset!
+							end
+							
+							@group.health_check!
+						else
+							break
+						end
+					else
+						break if check_timeout(child, state, age_clock, health_check_timeout, startup_timeout)
+					end
+				end
+				
+				return child.wait
+			end
+			
+			def wait_readable(child, timeout)
+				child.in.wait_readable(timeout)
+			rescue IO::TimeoutError
+				false
+			end
+			
+			def next_timeout(state, age_clock, health_check_timeout, startup_timeout)
+				return nil unless age_clock
+				
+				timeout = if state[:ready]
+					health_check_timeout
+				else
+					startup_timeout
+				end
+				
+				if timeout
+					remaining = timeout - age_clock.total
+					return 0 if remaining.negative?
+					
+					if interval = @group.health_check_interval
+						return [remaining, interval].min
+					else
+						return remaining
+					end
+				elsif interval = @group.health_check_interval
+					return interval
+				end
+			end
+			
+			def check_timeout(child, state, age_clock, health_check_timeout, startup_timeout)
+				return false unless age_clock
+				
+				if state[:ready]
+					if health_check_timeout && health_check_timeout < age_clock.total
+						health_check_failed(child, age_clock, health_check_timeout)
+						return true
+					end
+				else
+					if startup_timeout && startup_timeout < age_clock.total
+						startup_failed(child, age_clock, startup_timeout)
+						return true
+					end
+				end
+				
+				return false
+			end
+			
 			# Register the child (value) as running.
 			def insert(key, child)
 				if key
@@ -385,17 +452,6 @@ def delete(key, child)
 				@state.delete(child)
 			end
 
-			private
-			
-			if Fiber.respond_to?(:blocking?)
-				def fiber(&block)
-					Fiber.new(blocking: true, &block)
-				end
-			else
-				def fiber(&block)
-					Fiber.new(&block)
-				end
-			end
 		end
 	end
 end