client-libraries/ruby/lib/better_timeout.rb

   1 #--
   2 # = timeout.rb
   3 #
   4 # execution timeout
   5 #
   6 # = Copyright
   7 #
   8 # Copyright - (C) 2008  Evan Phoenix
   9 # Copyright:: (C) 2000  Network Applied Communication Laboratory, Inc.
  10 # Copyright:: (C) 2000  Information-technology Promotion Agency, Japan
  11 #
  12 #++
  13 #
  14 # = Description
  15 #
  16 # A way of performing a potentially long-running operation in a thread, and
  17 # terminating it's execution if it hasn't finished within fixed amount of
  18 # time.
  19 #
  20 # Previous versions of timeout didn't use a module for namespace. This version
  21 # provides both Timeout.timeout, and a backwards-compatible #timeout.
  22 #
  23 # = Synopsis
  24 #
  25 #   require 'timeout'
  26 #   status = Timeout::timeout(5) {
  27 #     # Something that should be interrupted if it takes too much time...
  28 #   }
  29 #
  30
  31 require 'thread'
  32
  33 module Timeout
  34
  35   ##
  36   # Raised by Timeout#timeout when the block times out.
  37
  38   class Error<Interrupt
  39   end
  40
  41   # A mutex to protect @requests
  42   @mutex = Mutex.new
  43
  44   # All the outstanding TimeoutRequests
  45   @requests = []
  46
  47   # Represents +thr+ asking for it to be timeout at in +secs+
  48   # seconds. At timeout, raise +exc+.
  49   class TimeoutRequest
  50     def initialize(secs, thr, exc)
  51       @left = secs
  52       @thread = thr
  53       @exception = exc
  54     end
  55
  56     attr_reader :thread, :left
  57
  58     # Called because +time+ seconds have gone by. Returns
  59     # true if the request has no more time left to run.
  60     def elapsed(time)
  61       @left -= time
  62       @left <= 0
  63     end
  64
  65     # Raise @exception if @thread.
  66     def cancel
  67       if @thread and @thread.alive?
  68         @thread.raise @exception, "execution expired"
  69       end
  70
  71       @left = 0
  72     end
  73
  74     # Abort this request, ie, we don't care about tracking
  75     # the thread anymore.
  76     def abort
  77       @thread = nil
  78       @left = 0
  79     end
  80   end
  81
  82   def self.add_timeout(time, exc)
  83
  84     @controller ||= Thread.new do
  85       while true
  86         if @requests.empty?
  87           sleep
  88           next
  89         end
  90
  91         min = nil
  92
  93         @mutex.synchronize do
  94           min = @requests.min { |a,b| a.left <=> b.left }
  95         end
  96
  97         slept_for = sleep(min.left)
  98
  99         @mutex.synchronize do
 100           @requests.delete_if do |r|
 101             if r.elapsed(slept_for)
 102               r.cancel
 103               true
 104             else
 105               false
 106             end
 107           end
 108         end
 109
 110       end
 111     end
 112
 113     req = TimeoutRequest.new(time, Thread.current, exc)
 114
 115     @mutex.synchronize do
 116       @requests << req
 117     end
 118
 119     @controller.run
 120
 121     return req
 122   end
 123
 124   ##
 125   # Executes the method's block. If the block execution terminates before +sec+
 126   # seconds has passed, it returns true. If not, it terminates the execution
 127   # and raises +exception+ (which defaults to Timeout::Error).
 128   #
 129   # Note that this is both a method of module Timeout, so you can 'include
 130   # Timeout' into your classes so they have a #timeout method, as well as a
 131   # module method, so you can call it directly as Timeout.timeout().
 132
 133   def timeout(sec, exception=Error)
 134     return yield if sec == nil or sec.zero?
 135     raise ThreadError, "timeout within critical session" if Thread.critical
 136
 137     req = Timeout.add_timeout sec, exception
 138
 139     begin
 140       yield sec
 141     ensure
 142       req.abort
 143     end
 144   end
 145
 146   module_function :timeout
 147
 148 end
 149
 150 ##
 151 # Identical to:
 152 #
 153 #   Timeout::timeout(n, e, &block).
 154 #
 155 # Defined for backwards compatibility with earlier versions of timeout.rb, see
 156 # Timeout#timeout.
 157
 158 def timeout(n, e=Timeout::Error, &block) # :nodoc:
 159   Timeout::timeout(n, e, &block)
 160 end
 161
 162 ##
 163 # Another name for Timeout::Error, defined for backwards compatibility with
 164 # earlier versions of timeout.rb.
 165
 166 TimeoutError = Timeout::Error # :nodoc:
 167
 168 if __FILE__ == $0
 169   p timeout(5) {
 170     45
 171   }
 172   p timeout(5, TimeoutError) {
 173     45
 174   }
 175   p timeout(nil) {
 176     54
 177   }
 178   p timeout(0) {
 179     54
 180   }
 181   p timeout(5) {
 182     loop {
 183       p 10
 184       sleep 1
 185     }
 186   }
 187 end
 188