From 2fcb2ca664836a06a27add697fb588801fd57d52 Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Wed, 3 Jan 2024 14:18:17 +1300
Subject: [PATCH] Allow Interrupt to retry the run loop after issuing #stop.
 (#297)

* Add Ruby v3.1 to coverage workflow.
---
 .github/workflows/coverage.yaml |  1 +
 lib/async/scheduler.rb          | 23 ++++++++++++--------
 test/async/scheduler.rb         | 37 +++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/coverage.yaml b/.github/workflows/coverage.yaml
index 41f6f584..c18f6f0c 100644
--- a/.github/workflows/coverage.yaml
+++ b/.github/workflows/coverage.yaml
@@ -20,6 +20,7 @@ jobs:
           - ubuntu
         
         ruby:
+          - "3.1"
           - "3.2"
           - "ruby-head"
     
diff --git a/lib/async/scheduler.rb b/lib/async/scheduler.rb
index ec38fb26..accf41e9 100644
--- a/lib/async/scheduler.rb
+++ b/lib/async/scheduler.rb
@@ -306,17 +306,22 @@ def run(...)
 			
 			initial_task = self.async(...) if block_given?
 			
-			# 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
-				while true
-					# If we are interrupted, we need to exit:
-					break if self.interrupted?
-					
-					# If we are finished, we need to exit:
-					break unless self.run_once
+			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
+					while true
+						# If we are interrupted, we need to exit:
+						break if self.interrupted?
+						
+						# If we are finished, we need to exit:
+						break unless self.run_once
+					end
 				end
+			rescue Interrupt
+				self.stop
+				retry
 			end
-			
+				
 			return initial_task
 		ensure
 			Console.logger.debug(self) {"Exiting run-loop because #{$! ? $! : 'finished'}."}
diff --git a/test/async/scheduler.rb b/test/async/scheduler.rb
index 1a34e910..e1285de1 100644
--- a/test/async/scheduler.rb
+++ b/test/async/scheduler.rb
@@ -87,6 +87,43 @@
 			scheduler.close
 			scheduler.interrupt
 		end
+		
+		it "can interrupt a scheduler from a different thread" do
+			finished = false
+			sleeping = Thread::Queue.new
+			
+			thread = Thread.new do
+				scheduler = Async::Scheduler.new
+				Fiber.set_scheduler(scheduler)
+				
+				scheduler.run do |task|
+					sleeping.push(true)
+					sleep
+				ensure
+					begin
+						sleeping.push(true)
+						sleep
+					ensure
+						finished = true
+					end
+				end
+			# rescue Interrupt
+			# 	# Ignore.
+			end
+			
+			expect(sleeping.pop).to be == true
+			expect(finished).to be == false
+			
+			thread.raise(Interrupt)
+			
+			expect(sleeping.pop).to be == true
+			expect(finished).to be == false
+			
+			thread.raise(Interrupt)
+			thread.join
+			
+			expect(finished).to be == true
+		end
 	end
 	
 	with '#block' do