From ca6397457922e574b1bc8dd49f2da591d2efc8a9 Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Sat, 24 Feb 2024 20:20:27 +1300
Subject: [PATCH] Better handling of interrupts during stop/terminate. (#308)

---
 lib/async/scheduler.rb  | 14 ++++++++++++--
 lib/async/task.rb       |  4 ++--
 test/async/scheduler.rb | 39 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/lib/async/scheduler.rb b/lib/async/scheduler.rb
index 75e18c1c..63e1245c 100644
--- a/lib/async/scheduler.rb
+++ b/lib/async/scheduler.rb
@@ -49,6 +49,13 @@ def scheduler_close
 			self.close
 		end
 		
+		# Terminate the scheduler. We deliberately ignore interrupts here, as this code can be called from an interrupt, and we don't want to be interrupted while cleaning up.
+		def terminate
+			Thread.handle_interrupt(::Interrupt => :never) do
+				super
+			end
+		end
+		
 		# @public Since `stable-v1`.
 		def close
 			# It's critical to stop all tasks. Otherwise they might be holding on to resources which are never closed/released correctly.
@@ -308,7 +315,7 @@ def run(...)
 			
 			begin
 				# In theory, we could use Exception here to be a little bit safer, but we've only shown the case for SignalException to be a problem, so let's not over-engineer this.
-				Thread.handle_interrupt(SignalException => :never) do
+				Thread.handle_interrupt(::SignalException => :never) do
 					while true
 						# If we are interrupted, we need to exit:
 						break if self.interrupted?
@@ -318,7 +325,10 @@ def run(...)
 					end
 				end
 			rescue Interrupt
-				self.stop
+				Thread.handle_interrupt(::SignalException => :never) do
+					self.stop
+				end
+				
 				retry
 			end
 				
diff --git a/lib/async/task.rb b/lib/async/task.rb
index d98b45fb..ecda7230 100644
--- a/lib/async/task.rb
+++ b/lib/async/task.rb
@@ -209,7 +209,7 @@ def wait
 		def stop(later = false)
 			if self.stopped?
 				# If we already stopped this task... don't try to stop it again:
-				return
+				return stopped!
 			end
 			
 			# If the fiber is alive, we need to stop it:
@@ -304,7 +304,7 @@ def stopped!
 			stopped = false
 			
 			begin
-				# We are bnot running, but children might be so we should stop them:
+				# We are not running, but children might be so we should stop them:
 				stop_children(true)
 			rescue Stop
 				stopped = true
diff --git a/test/async/scheduler.rb b/test/async/scheduler.rb
index effec312..126bdede 100644
--- a/test/async/scheduler.rb
+++ b/test/async/scheduler.rb
@@ -124,6 +124,45 @@
 			
 			expect(finished).to be == true
 		end
+		
+		it "ignores interrupts during termination" do
+			sleeping = Thread::Queue.new
+			
+			thread = Thread.new do
+				Thread.current.report_on_exception = false
+				
+				scheduler = Async::Scheduler.new
+				Fiber.set_scheduler(scheduler)
+				
+				scheduler.run do |task|
+					2.times do
+						task.async do
+							sleeping.push(true)
+							sleep
+						ensure
+							sleeping.push(true)
+							sleep
+						end
+					end
+				end
+			end
+			
+			# The first interrupt stops the tasks normally, but they enter sleep again:
+			expect(sleeping.pop).to be == true
+			thread.raise(Interrupt)
+			
+			# The second stop forcefully stops the two children tasks of the selector:
+			expect(sleeping.pop).to be == true
+			thread.raise(Interrupt)
+			
+			# The thread should now exit:
+			begin
+				thread.join
+			rescue Interrupt
+				# Ignore - this may happen:
+				# https://github.com/ruby/ruby/pull/10039
+			end
+		end
 	end
 	
 	with '#block' do