Skip to content

Commit

Permalink
Add hook for Timeout.timeout.
Browse files Browse the repository at this point in the history
  • Loading branch information
ioquatix committed Mar 30, 2021
1 parent 93753d7 commit 4c53dc9
Show file tree
Hide file tree
Showing 5 changed files with 71 additions and 5 deletions.
2 changes: 2 additions & 0 deletions include/ruby/fiber/scheduler.h
Expand Up @@ -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);

Expand Down
8 changes: 7 additions & 1 deletion lib/timeout.rb
Expand Up @@ -76,9 +76,15 @@ def exception(*)
# 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|
Expand Down
8 changes: 8 additions & 0 deletions scheduler.c
Expand Up @@ -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;

Expand All @@ -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");

Expand Down Expand Up @@ -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)
{
Expand Down
28 changes: 24 additions & 4 deletions test/fiber/scheduler.rb
Expand Up @@ -81,10 +81,12 @@ def run
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
Expand Down Expand Up @@ -127,6 +129,24 @@ def current_time
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

Expand Down
30 changes: 30 additions & 0 deletions 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

0 comments on commit 4c53dc9

Please sign in to comment.