diff options
| -rw-r--r-- | include/ruby/fiber/scheduler.h | 2 | ||||
| -rw-r--r-- | lib/timeout.rb | 8 | ||||
| -rw-r--r-- | scheduler.c | 8 | ||||
| -rw-r--r-- | test/fiber/scheduler.rb | 28 | ||||
| -rw-r--r-- | test/fiber/test_timeout.rb | 30 |
5 files changed, 71 insertions, 5 deletions
diff --git a/include/ruby/fiber/scheduler.h b/include/ruby/fiber/scheduler.h index 36d1726b19..213a8b8a1f 100644 --- a/include/ruby/fiber/scheduler.h +++ b/include/ruby/fiber/scheduler.h @@ -22,6 +22,8 @@ VALUE rb_fiber_scheduler_make_timeout(struct timeval *timeout); VALUE rb_fiber_scheduler_close(VALUE scheduler); +VALUE rb_fiber_scheduler_timeout_raise(VALUE scheduler, VALUE duration); + VALUE rb_fiber_scheduler_kernel_sleep(VALUE scheduler, VALUE duration); VALUE rb_fiber_scheduler_kernel_sleepv(VALUE scheduler, int argc, VALUE * argv); 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| diff --git a/scheduler.c b/scheduler.c index f2b1b00fa1..3403eb1801 100644 --- a/scheduler.c +++ b/scheduler.c @@ -17,6 +17,7 @@ static ID id_close; static ID id_block; static ID id_unblock; +static ID id_timeout_raise; static ID id_kernel_sleep; static ID id_process_wait; @@ -32,6 +33,7 @@ Init_Fiber_Scheduler(void) id_block = rb_intern_const("block"); id_unblock = rb_intern_const("unblock"); + id_timeout_raise = rb_intern_const("timeout_raise"); id_kernel_sleep = rb_intern_const("kernel_sleep"); id_process_wait = rb_intern_const("process_wait"); @@ -109,6 +111,12 @@ rb_fiber_scheduler_make_timeout(struct timeval *timeout) } VALUE +rb_fiber_scheduler_timeout_raise(VALUE scheduler, VALUE timeout) +{ + return rb_funcall(scheduler, id_timeout_raise, 1, timeout); +} + +VALUE rb_fiber_scheduler_kernel_sleep(VALUE scheduler, VALUE timeout) { return rb_funcall(scheduler, id_kernel_sleep, 1, timeout); diff --git a/test/fiber/scheduler.rb b/test/fiber/scheduler.rb index f2fb304e19..8a8585fcbe 100644 --- a/test/fiber/scheduler.rb +++ b/test/fiber/scheduler.rb @@ -81,10 +81,12 @@ class Scheduler waiting, @waiting = @waiting, {} waiting.each do |fiber, timeout| - if timeout <= time - fiber.resume - else - @waiting[fiber] = timeout + if fiber.alive? + if timeout <= time + fiber.resume + else + @waiting[fiber] = timeout + end end end end @@ -127,6 +129,24 @@ class Scheduler Process.clock_gettime(Process::CLOCK_MONOTONIC) end + def timeout_raise(duration, klass, message, &block) + fiber = Fiber.current + + self.fiber do + sleep(duration) + + if fiber&.alive? + fiber.raise(klass, message) + end + end + + begin + yield(duration) + ensure + fiber = nil + end + end + def process_wait(pid, flags) # $stderr.puts [__method__, pid, flags, Fiber.current].inspect diff --git a/test/fiber/test_timeout.rb b/test/fiber/test_timeout.rb new file mode 100644 index 0000000000..b974aa0e35 --- /dev/null +++ b/test/fiber/test_timeout.rb @@ -0,0 +1,30 @@ +# frozen_string_literal: true +require 'test/unit' +require_relative 'scheduler' + +require 'timeout' + +class TestFiberTimeout < Test::Unit::TestCase + def test_timeout_raise + error = nil + + thread = Thread.new do + scheduler = Scheduler.new + Fiber.set_scheduler scheduler + + Fiber.schedule do + begin + Timeout.timeout(0.01) do + sleep(1) + end + rescue + error = $! + end + end + end + + thread.join + + assert_kind_of(Timeout::Error, error) + end +end |