Add more flexible timeout.

Author: ioquatix

lib/async/scheduler.rb

@@ -7,6 +7,7 @@
 
 require_relative "clock"
 require_relative "task"
+require_relative "timeout"
 require_relative "worker_pool"
 
 require "io/event"
@@ -539,7 +540,7 @@ def fiber(...)
 		# @parameter duration [Numeric] The time in seconds, in which the task should complete.
 		# @parameter exception [Class] The exception class to raise.
 		# @parameter message [String] The message to pass to the exception.
-		# @yields {|duration| ...} The block to execute with a timeout.
+		# @yields {|timeout| ...} The block to execute with a timeout.
 		def with_timeout(duration, exception = TimeoutError, message = "execution expired", &block)
 			fiber = Fiber.current
 
@@ -549,7 +550,11 @@ def with_timeout(duration, exception = TimeoutError, message = "execution expire
 				end
 			end
 
-			yield timer
+			if block.arity.zero?
+				yield
+			else
+				yield Timeout.new(@timers, timer, duration)
+			end
 		ensure
 			timer&.cancel!
 		end
@@ -564,7 +569,7 @@ def with_timeout(duration, exception = TimeoutError, message = "execution expire
 		# @parameter message [String] The message to pass to the exception.
 		# @yields {|duration| ...} The block to execute with a timeout.
 		def timeout_after(duration, exception, message, &block)
-			with_timeout(duration, exception, message) do |timer|
+			with_timeout(duration, exception, message) do
 				yield duration
 			end
 		end

lib/async/timeout.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+# Released under the MIT License.
+# Copyright, 2025, by Samuel Williams.
+
+module Async
+	# Represents a flexible timeout that can be rescheduled or extended.
+	# @public Since *Async v2.24*.
+	class Timeout
+		# Initialize a new timeout.
+		def initialize(timers, handle, duration = nil)
+			@timers = timers
+			@handle = handle
+			@duration = duration || (handle.time - timers.now)
+		end
+		
+		# @attribute [Numeric] The duration of the timeout.
+		attr :duration
+		
+		# Update the duration of the timeout, rescheduling it if necessary.
+		#
+		# The duration is relative to the time the timeout was created.
+		#
+		# @parameter value [Numeric] The new duration to assign to the timeout.
+		def duration=(value)
+			delta = value - @duration
+			self.reschedule(time + delta, value)
+		end
+		
+		# Adjust the timeout by the specified duration, rescheduling it if necessary.
+		#
+		# @parameter duration [Numeric] The duration to adjust the timeout by.
+		# @returns [Numeric] The new time at which the timeout will occur.
+		def adjust(duration)
+			self.reschedule(time + duration, @duration + duration)
+		end
+		
+		# @returns [Numeric] The time at which the timeout will occur.
+		def time
+			@handle.time
+		end
+		
+		# Assign a new time to the timeout, rescheduling it if necessary.
+		#
+		# @parameter value [Numeric] The new time to assign to the timeout.
+		# @returns [Numeric] The new time at which the timeout will occur.
+		def time=(value)
+			self.reschedule(value)
+		end
+		
+		# Cancel the timeout.
+		def cancel!
+			@handle.cancel!
+		end
+		
+		# @returns [Boolean] Whether the timeout has been cancelled.
+		def cancelled?
+			@handle.cancelled?
+		end
+		
+		# Reschedule the timeout to occur at the specified time.
+		#
+		# @parameter time [Numeric] The new time to schedule the timeout for.
+		# @parameter duration [Numeric | Nil] The new duration to assign to the timeout.
+		# @returns [Numeric] The new time at which the timeout will occur.
+		private def reschedule(time, duration = nil)
+			if block = @handle&.block
+				@handle.cancel!
+				
+				@duration = duration || (time - @timers.now)
+				@handle = @timers.schedule(time, block)
+				
+				return time
+			else
+				raise RuntimeError, "Cannot reschedule a cancelled timeout!"
+			end
+		end
+	end
+end

test/async/timeout.rb

@@ -0,0 +1,48 @@
+require "sus/fixtures/async"
+
+describe Async::Timeout do
+	include Sus::Fixtures::Async::ReactorContext
+	
+	it "can schedule a timeout" do
+		scheduler.with_timeout(1) do |timeout|
+			expect(timeout.time).to be >= 0
+			expect(timeout.duration).to be == 1
+		end
+	end
+	
+	with "#adjust" do
+		it "can adjust the timeout" do
+			scheduler.with_timeout(1) do |timeout|
+				timeout.adjust(1)
+				expect(timeout.duration).to be == 2
+			end
+		end
+	end
+	
+	with "#duration=" do
+	 it "can set the timeout duration" do
+			scheduler.with_timeout(1) do |timeout|
+				timeout.duration = 2
+				expect(timeout.duration).to be == 2
+			end
+		end
+	end
+	
+	with "#time=" do
+		it "can set the timeout time" do
+			scheduler.with_timeout(1) do |timeout|
+				timeout.time = timeout.time + 1
+				expect(timeout.duration).to (be > 1).and(be < 2)
+			end
+		end
+	end
+	
+	with "#cancel!" do
+		it "can cancel the timeout" do
+			scheduler.with_timeout(1) do |timeout|
+				timeout.cancel!
+				expect(timeout).to be(:cancelled?)
+			end
+		end
+	end
+end