Reimplement Timeout.timeout with a single thread and a Queue

Author: eregon

lib/timeout.rb

@@ -1,4 +1,4 @@
-# frozen_string_literal: false
+# frozen_string_literal: true
 # Timeout long-running blocks
 #
 # == Synopsis
@@ -23,7 +23,7 @@
 # Copyright:: (C) 2000  Information-technology Promotion Agency, Japan
 
 module Timeout
-  VERSION = "0.2.0".freeze
+  VERSION = "0.2.0"
 
   # Raised by Timeout.timeout when the block times out.
   class Error < RuntimeError
@@ -50,9 +50,79 @@ def exception(*)
   end
 
   # :stopdoc:
-  THIS_FILE = /\A#{Regexp.quote(__FILE__)}:/o
-  CALLER_OFFSET = ((c = caller[0]) && THIS_FILE =~ c) ? 1 : 0
-  private_constant :THIS_FILE, :CALLER_OFFSET
+  CONDVAR = ConditionVariable.new
+  QUEUE = Queue.new
+  QUEUE_MUTEX = Mutex.new
+  TIMEOUT_THREAD_MUTEX = Mutex.new
+  @timeout_thread = nil
+  private_constant :CONDVAR, :QUEUE, :QUEUE_MUTEX, :TIMEOUT_THREAD_MUTEX
+
+  class Request
+    attr_reader :deadline
+
+    def initialize(thread, timeout, exception_class, message)
+      @thread = thread
+      @deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
+      @exception_class = exception_class
+      @message = message
+
+      @mutex = Mutex.new
+      @done = false
+    end
+
+    def done?
+      @done
+    end
+
+    def expired?(now)
+      now >= @deadline and !done?
+    end
+
+    def interrupt
+      @mutex.synchronize do
+        unless @done
+          @thread.raise @exception_class, @message
+          @done = true
+        end
+      end
+    end
+
+    def finished
+      @mutex.synchronize do
+        @done = true
+      end
+    end
+  end
+  private_constant :Request
+
+  def self.ensure_timeout_thread_created
+    unless @timeout_thread
+      TIMEOUT_THREAD_MUTEX.synchronize do
+        @timeout_thread ||= Thread.new do
+          requests = []
+          while true
+            until QUEUE.empty? and !requests.empty? # wait to have at least one request
+              req = QUEUE.pop
+              requests << req unless req.done?
+            end
+            closest_deadline = requests.min_by(&:deadline).deadline
+
+            now = 0.0
+            QUEUE_MUTEX.synchronize do
+              while (now = Process.clock_gettime(Process::CLOCK_MONOTONIC)) < closest_deadline and QUEUE.empty?
+                CONDVAR.wait(QUEUE_MUTEX, closest_deadline - now)
+              end
+            end
+
+            requests.each do |req|
+              req.interrupt if req.expired?(now)
+            end
+            requests.reject!(&:done?)
+          end
+        end
+      end
+    end
+  end
   # :startdoc:
 
   # Perform an operation in a block, raising an error if it takes longer than
@@ -83,51 +153,32 @@ def exception(*)
   def timeout(sec, klass = nil, message = nil, &block)   #:yield: +sec+
     return yield(sec) if sec == nil or sec.zero?
 
-    message ||= "execution expired".freeze
+    message ||= "execution expired"
 
     if Fiber.respond_to?(:current_scheduler) && (scheduler = Fiber.current_scheduler)&.respond_to?(:timeout_after)
       return scheduler.timeout_after(sec, klass || Error, message, &block)
     end
 
-    from = "from #{caller_locations(1, 1)[0]}" if $DEBUG
-    e = Error
-    bl = proc do |exception|
+    Timeout.ensure_timeout_thread_created
+    perform = Proc.new do |exc|
+      request = Request.new(Thread.current, sec, exc, message)
+      QUEUE_MUTEX.synchronize do
+        QUEUE << request
+        CONDVAR.signal
+      end
       begin
-        x = Thread.current
-        y = Thread.start {
-          Thread.current.name = from
-          begin
-            sleep sec
-          rescue => e
-            x.raise e
-          else
-            x.raise exception, message
-          end
-        }
         return yield(sec)
       ensure
-        if y
-          y.kill
-          y.join # make sure y is dead.
-        end
+        request.finished
       end
     end
+
     if klass
-      begin
-        bl.call(klass)
-      rescue klass => e
-        message = e.message
-        bt = e.backtrace
-      end
+      perform.call(klass)
     else
-      bt = Error.catch(message, &bl)
+      backtrace = Error.catch(&perform)
+      raise Error, message, backtrace
     end
-    level = -caller(CALLER_OFFSET).size-2
-    while THIS_FILE =~ bt[level]
-      bt.delete_at(level)
-    end
-    raise(e, message, bt)
   end
-
   module_function :timeout
 end

test/test_timeout.rb

@@ -10,6 +10,18 @@ def test_non_timing_out_code_is_successful
     end
   end
 
+  def test_included
+    c = Class.new do
+      include Timeout
+      def test
+        timeout(1) { :ok }
+      end
+    end
+    assert_nothing_raised do
+      assert_equal :ok, c.new.test
+    end
+  end
+
   def test_yield_param
     assert_equal [5, :ok], Timeout.timeout(5){|s| [s, :ok] }
   end
@@ -43,6 +55,7 @@ def test_skip_rescue
         begin
           sleep 3
         rescue Exception => e
+          flunk "should not see any exception but saw #{e.inspect}"
         end
       end
     end