Add hook for `Timeout.timeout`.

1 files changed, 7 insertions, 1 deletions

diff --git a/lib/timeout.rb b/lib/timeout.rb
index 9026ad51d6..43f26f5869 100644
--- a/lib/timeout.rb
+++ b/lib/timeout.rb
@@ -76,9 +76,15 @@ module Timeout
   # Note that this is both a method of module Timeout, so you can <tt>include
   # Timeout</tt> into your classes so they have a #timeout method, as well as
   # a module method, so you can call it directly as Timeout.timeout().
-  def timeout(sec, klass = nil, message = nil)   #:yield: +sec+
+  def timeout(sec, klass = nil, message = nil, &block)   #:yield: +sec+
     return yield(sec) if sec == nil or sec.zero?
+
     message ||= "execution expired".freeze
+
+    if scheduler = Fiber.scheduler and scheduler.respond_to?(:timeout_raise)
+      return scheduler.timeout_raise(sec, klass || Error, message, &block)
+    end
+
     from = "from #{caller_locations(1, 1)[0]}" if $DEBUG
     e = Error
     bl = proc do |exception|