From 61f26882580c41735a4fee47aeb4826e5ab77d5e Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Sun, 22 Dec 2024 14:59:43 +1300
Subject: [PATCH] More robust signal handling.

---
 lib/async/container/controller.rb | 61 ++++++++++++++++++-------------
 lib/async/container/group.rb      | 14 +++++--
 lib/async/container/process.rb    |  3 +-
 releases.md                       |  1 +
 4 files changed, 49 insertions(+), 30 deletions(-)

diff --git a/lib/async/container/controller.rb b/lib/async/container/controller.rb
index bc6da05..fb7b3a7 100644
--- a/lib/async/container/controller.rb
+++ b/lib/async/container/controller.rb
@@ -187,47 +187,56 @@ def reload
 			def run
 				@notify&.status!("Initializing...")
 				
+				with_signal_handlers do
+					self.start
+					
+					while @container&.running?
+						begin
+							@container.wait
+						rescue SignalException => exception
+							if handler = @signals[exception.signo]
+								begin
+									handler.call
+								rescue SetupError => error
+									Console.error(self) {error}
+								end
+							else
+								raise
+							end
+						end
+					end
+				rescue Interrupt
+					self.stop
+				rescue Terminate
+					self.stop(false)
+				ensure
+					self.stop(false)
+				end
+			end
+			
+			private def with_signal_handlers
 				# I thought this was the default... but it doesn't always raise an exception unless you do this explicitly.
-				# We use `Thread.current.raise(...)` so that exceptions are filtered through `Thread.handle_interrupt` correctly.
+				
 				interrupt_action = Signal.trap(:INT) do
-					# $stderr.puts "Received INT signal, terminating...", caller
+					# We use `Thread.current.raise(...)` so that exceptions are filtered through `Thread.handle_interrupt` correctly.
+					# $stderr.puts "Received Interrupt signal, terminating...", caller
 					::Thread.current.raise(Interrupt)
 				end
 				
 				terminate_action = Signal.trap(:TERM) do
-					# $stderr.puts "Received TERM signal, terminating...", caller
+					# $stderr.puts "Received Terminate signal, terminating...", caller
 					::Thread.current.raise(Terminate)
 				end
 				
 				hangup_action = Signal.trap(:HUP) do
-					# $stderr.puts "Received HUP signal, restarting...", caller
+					# $stderr.puts "Received Hangup signal, restarting...", caller
 					::Thread.current.raise(Hangup)
 				end
 				
-				self.start
-				
-				while @container&.running?
-					begin
-						@container.wait
-					rescue SignalException => exception
-						if handler = @signals[exception.signo]
-							begin
-								handler.call
-							rescue SetupError => error
-								Console.error(self) {error}
-							end
-						else
-							raise
-						end
-					end
+				::Thread.handle_interrupt(SignalException => :on_blocking) do
+					yield
 				end
-			rescue Interrupt
-				self.stop
-			rescue Terminate
-				self.stop(false)
 			ensure
-				self.stop(false)
-				
 				# Restore the interrupt handler:
 				Signal.trap(:INT, interrupt_action)
 				Signal.trap(:TERM, terminate_action)
diff --git a/lib/async/container/group.rb b/lib/async/container/group.rb
index 91d47dd..82a844a 100644
--- a/lib/async/container/group.rb
+++ b/lib/async/container/group.rb
@@ -142,11 +142,19 @@ def wait_for_children(duration = nil)
 				
 				if !@running.empty?
 					# Maybe consider using a proper event loop here:
+					if ready = self.select(duration)
+						ready.each do |io|
+							@running[io].resume
+						end
+					end
+				end
+			end
+			
+			def select(duration)
+				::Thread.handle_interrupt(SignalException => :immediate) do
 					readable, _, _ = ::IO.select(@running.keys, nil, nil, duration)
 					
-					readable&.each do |io|
-						@running[io].resume
-					end
+					return readable
 				end
 			end
 			
diff --git a/lib/async/container/process.rb b/lib/async/container/process.rb
index c14e988..cb0574c 100644
--- a/lib/async/container/process.rb
+++ b/lib/async/container/process.rb
@@ -69,7 +69,8 @@ def self.fork(**options)
 						Signal.trap(:INT) {::Thread.current.raise(Interrupt)}
 						Signal.trap(:TERM) {::Thread.current.raise(Terminate)}
 						
-						begin
+						# This could be a configuration option:
+						::Thread.handle_interrupt(SignalException => :immediate) do
 							yield Instance.for(process)
 						rescue Interrupt
 							# Graceful exit.
diff --git a/releases.md b/releases.md
index 085fbef..622e307 100644
--- a/releases.md
+++ b/releases.md
@@ -2,4 +2,5 @@
 
 ## Unreleased
 
+  - Improve container signal handling reliability by using `Thread.handle_interrupt` except at known safe points.
   - Improved logging when child process fails and container startup.