Add hook for `Timeout.timeout`.

This commit is contained in:
Samuel Williams 2021-02-11 19:17:54 +13:00
parent 93753d7ee7
commit 4c53dc970b
Notes: git 2021-03-30 14:39:05 +09:00
5 changed files with 71 additions and 5 deletions

View File

@ -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);

View File

@ -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|

View File

@ -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");
@ -108,6 +110,12 @@ rb_fiber_scheduler_make_timeout(struct timeval *timeout)
return Qnil;
}
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)
{

View File

@ -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

View File

@ -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