Add more flexible timeout. (#386)

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)
+			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,88 @@
+# 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)
+			@timers = timers
+			@handle = handle
+		end
+		
+		# @returns [Numeric] The time remaining until the timeout occurs, in seconds.
+		def duration
+			@handle.time - @timers.now
+		end
+		
+		# Update the duration of the timeout.
+		#
+		# The duration is relative to the current time, e.g. setting the duration to 5 means the timeout will occur in 5 seconds from now.
+		#
+		# @parameter value [Numeric] The new duration to assign to the timeout, in seconds.
+		def duration=(value)
+			self.reschedule(@timers.now + value)
+		end
+		
+		# Adjust the timeout by the specified duration.
+		#
+		# The duration is relative to the timeout time, e.g. adjusting the timeout by 5 increases the current duration by 5 seconds.
+		#
+		# @parameter duration [Numeric] The duration to adjust the timeout by, in seconds.
+		# @returns [Numeric] The new time at which the timeout will occur.
+		def adjust(duration)
+			self.reschedule(time + duration)
+		end
+		
+		# @returns [Numeric] The time at which the timeout will occur, in seconds since {now}.
+		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
+		
+		# @returns [Numeric] The current time in the scheduler, relative to the time of this timeout, in seconds.
+		def now
+			@timers.now
+		end
+		
+		# Cancel the timeout, preventing it from executing.
+		def cancel!
+			@handle.cancel!
+		end
+		
+		# @returns [Boolean] Whether the timeout has been cancelled.
+		def cancelled?
+			@handle.cancelled?
+		end
+		
+		# Raised when attempting to reschedule a cancelled timeout.
+		class CancelledError < RuntimeError
+		end
+		
+		# Reschedule the timeout to occur at the specified time.
+		#
+		# @parameter time [Numeric] The new time to schedule the timeout for.
+		# @returns [Numeric] The new time at which the timeout will occur.
+		private def reschedule(time)
+			if block = @handle&.block
+				@handle.cancel!
+				
+				@handle = @timers.schedule(time, block)
+				
+				return time
+			else
+				raise CancelledError, "Cannot reschedule a cancelled timeout!"
+			end
+		end
+	end
+end

releases.md

@@ -5,6 +5,34 @@
   - Ruby v3.1 support is dropped.
   - `Async::Wrapper` which was previously deprecated, is now removed.
 
+### Flexible Timeouts
+
+When {ruby Async::Scheduler#with_timeout} is invoked with a block, it can receive a {ruby Async::Timeout} instance. This allows you to adjust or cancel the timeout while the block is executing. This is useful for long-running tasks that may need to adjust their timeout based on external factors.
+
+``` ruby
+Async do
+	Async::Scheduler.with_timeout(5) do |timeout|
+		# Do some work that may take a while...
+		
+		if some_condition
+			timeout.cancel! # Cancel the timeout
+		else
+			# Add 10 seconds to the current timeout:
+			timeout.adjust(10)
+			
+			# Reduce the timeout by 10 seconds:
+			timeout.adjust(-10)
+			
+			# Set the timeout to 10 seconds from now:
+			timeout.duration = 10
+			
+			# Increase the current duration:
+			timeout.duration += 10
+		end
+	end
+end
+```
+
 ## v2.23.0
 
   - Rename `ASYNC_SCHEDULER_DEFAULT_WORKER_POOL` to `ASYNC_SCHEDULER_WORKER_POOL`.

test/async/timeout.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+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 > 0).and(be <= 1)
+		end
+	end
+	
+	with "#now" do
+		it "can get the current time" do
+			scheduler.with_timeout(1) do |timeout|
+				expect(timeout.now).to be >= 0
+				expect(timeout.now).to be <= timeout.time
+			end
+		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 > 1).and(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 > 1).and(be <= 2)
+			end
+		end
+		
+		it "can increase the timeout duration" do
+			scheduler.with_timeout(1) do |timeout|
+				timeout.duration += 2
+				expect(timeout.duration).to (be > 2).and(be <= 3)
+			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
+		
+		it "can't reschedule a cancelled timeout" do
+			scheduler.with_timeout(1) do |timeout|
+				timeout.cancel!
+				expect do
+					timeout.adjust(1)
+				end.to raise_exception(Async::Timeout::CancelledError)
+			end
+		end
+	end
+end