lib/timeout/extensions.rb in timeout-extensions-0.0.0 vs lib/timeout/extensions.rb in timeout-extensions-0.1.0

@@ -1,45 +1,50 @@
-require 'timeout/extensions/version'
+require "timeout/extensions/version"
 
+# Core extensions to Thread
+#
+# Adds accessors to customize timeout and sleep within a thread,
+# if you have a better friendlier implementation of both methods,
+# or if your code breaks with the stdlib implementations.
+#
 class Thread
   attr_accessor :timeout_handler
   attr_accessor :sleep_handler
 end
 
 module Timeout::Extensions
-  def self.extended(mod)
-    mod.singleton_class.class_eval do
-      alias_method :timeout_without_handler, :timeout
-      alias_method :timeout, :timeout_with_handler
+  module TimeoutMethods
+    def timeout(*args, &block)
+      if (timeout_handler = Thread.current.timeout_handler)
+        timeout_handler.call(*args, &block)
+      else
+        super
+      end
     end
   end
 
-  def timeout_with_handler(*args, &block)
-    if timeout_handler = Thread.current.timeout_handler
-      timeout_handler.call(*args, &block)
-    else
-      timeout_without_handler(*args, &block)
+  module KernelMethods
+    def sleep(*args)
+      if (sleep_handler = Thread.current.sleep_handler)
+        sleep_handler.call(*args)
+      else
+        super
+      end
     end
   end
-  module_function :timeout_with_handler
-  public :timeout_with_handler
-end
 
-module Kernel::Extensions
-  def self.included(mod)
-    mod.class_eval do
-      alias_method :sleep_without_handler, :sleep
-      alias_method :sleep, :sleep_with_handler
-    end
+  # in order for prepend to work, I have to do it in the Timeout module singleton class
+  class << ::Timeout
+    prepend TimeoutMethods
   end
 
-  def sleep_with_handler(*args)
-    if sleep_handler = Thread.current.sleep_handler
-      sleep_handler.call(*args)
-    else
-      sleep_without_handler(*args)
-    end
+  # ditto for Kernel
+  class << ::Kernel
+    prepend KernelMethods
   end
-end
 
-Timeout.extend Timeout::Extensions
-include Kernel::Extensions
+  # this is an hack so that calling "sleep(2)" works. Amazingly, the message doesn't get
+  # sent to Kernel.sleep code path.
+  # https://bugs.ruby-lang.org/issues/12535
+  #
+  ::Object.prepend KernelMethods
+end