Fix handling of timeout accessing scheduler outside of non-blocking context.

4 files changed, 34 insertions, 1 deletions

diff --git a/cont.c b/cont.c
index 0d95a097b4..bdd308569c 100644
--- a/cont.c
+++ b/cont.c
@@ -1997,6 +1997,20 @@ rb_fiber_s_scheduler(VALUE klass)
 
 /*
  *  call-seq:
+ *     Fiber.current_scheduler -> obj or nil
+ *
+ *  Returns the Fiber scheduler, that was last set for the current thread with Fiber.set_scheduler
+ *  iff the current fiber is non-blocking.
+ *
+ */
+static VALUE
+rb_fiber_current_scheduler(VALUE klass)
+{
+    return rb_fiber_scheduler_current();
+}
+
+/*
+ *  call-seq:
  *     Fiber.set_scheduler(scheduler) -> scheduler
  *
  *  Sets the Fiber scheduler for the current thread. If the scheduler is set, non-blocking
@@ -3084,6 +3098,7 @@ Init_Cont(void)
     rb_define_singleton_method(rb_cFiber, "blocking?", rb_fiber_s_blocking_p, 0);
     rb_define_singleton_method(rb_cFiber, "scheduler", rb_fiber_s_scheduler, 0);
     rb_define_singleton_method(rb_cFiber, "set_scheduler", rb_fiber_set_scheduler, 1);
+    rb_define_singleton_method(rb_cFiber, "current_scheduler", rb_fiber_current_scheduler, 0);
 
     rb_define_singleton_method(rb_cFiber, "schedule", rb_fiber_s_schedule, -1);
 
diff --git a/lib/timeout.rb b/lib/timeout.rb
index dc8eb24a10..11db4be973 100644
--- a/lib/timeout.rb
+++ b/lib/timeout.rb
@@ -84,7 +84,7 @@ module Timeout
 
     message ||= "execution expired".freeze
 
-    if (scheduler = Fiber.scheduler)&.respond_to?(:timeout_after)
+    if (scheduler = Fiber.current_scheduler)&.respond_to?(:timeout_after)
       return scheduler.timeout_after(sec, klass || Error, message, &block)
     end
 
diff --git a/test/fiber/test_scheduler.rb b/test/fiber/test_scheduler.rb
index 72bde9fcc3..eeb0d67ec5 100644
--- a/test/fiber/test_scheduler.rb
+++ b/test/fiber/test_scheduler.rb
@@ -73,4 +73,20 @@ class TestFiberScheduler < Test::Unit::TestCase
 
     thread.join
   end
+
+  def test_current_scheduler
+    thread = Thread.new do
+      scheduler = Scheduler.new
+      Fiber.set_scheduler scheduler
+
+      assert Fiber.scheduler
+      refute Fiber.current_scheduler
+
+      Fiber.schedule do
+        assert Fiber.current_scheduler
+      end
+    end
+
+    thread.join
+  end
 end
diff --git a/test/fiber/test_timeout.rb b/test/fiber/test_timeout.rb
index c17092b8e6..5493606be0 100644
--- a/test/fiber/test_timeout.rb
+++ b/test/fiber/test_timeout.rb
@@ -37,6 +37,8 @@ class TestFiberTimeout < Test::Unit::TestCase
       scheduler = Scheduler.new
       Fiber.set_scheduler scheduler
 
+      assert_nil Fiber.current_scheduler
+
       Timeout.timeout(1) do
         message = MESSAGE
       end