From 176fca1abfe4e223908ac29bd7db95506a008456 Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Tue, 2 Jan 2024 13:49:10 +1300
Subject: [PATCH] Add support for IO#timeout in `io_read`, `io_write` and
 `io_wait`. (#296)

---
 lib/async/scheduler.rb | 43 +++++++++++++++++++++++++++++++++++++++---
 test/io.rb             | 30 +++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/lib/async/scheduler.rb b/lib/async/scheduler.rb
index 4efe82aa..ec38fb26 100644
--- a/lib/async/scheduler.rb
+++ b/lib/async/scheduler.rb
@@ -165,29 +165,66 @@ def address_resolve(hostname)
 			::Resolv.getaddresses(hostname)
 		end
 		
+		
+		if IO.method_defined?(:timeout)
+			private def get_timeout(io)
+				io.timeout
+			end
+		else
+			private def get_timeout(io)
+				nil
+			end
+		end
+		
 		# @asynchronous May be non-blocking..
 		def io_wait(io, events, timeout = nil)
 			fiber = Fiber.current
 			
 			if timeout
+				# If an explicit timeout is specified, we expect that the user will handle it themselves:
 				timer = @timers.after(timeout) do
 					fiber.transfer
 				end
+			elsif timeout = get_timeout(io)
+				# Otherwise, if we default to the io's timeout, we raise an exception:
+				timer = @timers.after(timeout) do
+					fiber.raise(::IO::TimeoutError, "Timeout while waiting for IO to become ready!")
+				end
 			end
 			
 			return @selector.io_wait(fiber, io, events)
 		ensure
 			timer&.cancel
 		end
-
+		
 		if ::IO::Event::Support.buffer?
 			def io_read(io, buffer, length, offset = 0)
-				@selector.io_read(Fiber.current, io, buffer, length, offset)
+				fiber = Fiber.current
+				
+				if timeout = get_timeout(io)
+					timer = @timers.after(timeout) do
+						fiber.raise(::IO::TimeoutError, "execution expired")
+					end
+				end
+				
+				@selector.io_read(fiber, io, buffer, length, offset)
+			ensure
+				timer&.cancel
 			end
 			
 			if RUBY_ENGINE != "ruby" || RUBY_VERSION >= "3.3.0"
 				def io_write(io, buffer, length, offset = 0)
-					@selector.io_write(Fiber.current, io, buffer, length, offset)
+					fiber = Fiber.current
+				
+					if timeout = get_timeout(io)
+						timer = @timers.after(timeout) do
+							fiber.raise(::IO::TimeoutError, "execution expired")
+						end
+					end
+					
+					@selector.io_write(fiber, io, buffer, length, offset)
+				ensure
+					timer&.cancel
 				end
 			end
 		end
diff --git a/test/io.rb b/test/io.rb
index f73e938c..ce32a43c 100644
--- a/test/io.rb
+++ b/test/io.rb
@@ -30,5 +30,35 @@
 			input.close
 			output.close
 		end
+		
+		it "can read with timeout" do
+			skip_unless_constant_defined(:TimeoutError, IO)
+			
+			input, output = IO.pipe
+			input.timeout = 0.001
+			
+			expect do
+				line = input.gets
+			end.to raise_exception(::IO::TimeoutError)
+		end
+		
+		it "can wait readable with default timeout" do
+			skip_unless_constant_defined(:TimeoutError, IO)
+			
+			input, output = IO.pipe
+			input.timeout = 0.001
+			
+			expect do
+				# This behaviour is not consistent with non-fiber scheduler IO.
+				# However, this is the best we can do without fixing CRuby.
+				input.wait_readable
+			end.to raise_exception(::IO::TimeoutError)
+		end
+		
+		it "can wait readable with explicit timeout" do
+			input, output = IO.pipe
+			
+			expect(input.wait_readable(0)).to be_nil
+		end
 	end
 end