lib/double_restraint.rb in double_restraint-1.0.0 vs lib/double_restraint.rb in double_restraint-1.0.1

@@ -1,10 +1,12 @@
 # frozen_string_literal: true
 
 require "restrainer"
 
 class DoubleRestraint
+  attr_reader :timeout, :long_running_timeout
+
   # @param name [String, Symbol] The name of the restraint.
   # @param timeout [Numeric] The first timeout that will be yielded to the block.
   # @param long_running_timeout [Numeric] The timeout that will be yielded to the block if
   #        the block times out the first time it is excuted.
   # @param long_running_limit [Integer] The maximum number of times the block can be run
@@ -18,26 +20,24 @@
   #        If this is not set, the default value set for `Restrainer.redis` will be used.
   def initialize(name, timeout:, long_running_timeout:, long_running_limit:, limit: nil, timeout_errors: [Timeout::Error], redis: nil)
     @timeout = timeout
     @long_running_timeout = long_running_timeout
     @timeout_errors = Array(timeout_errors)
-    @restrainer = Restrainer.new("DoubleRestrainer(#{name})", limit: limit, redis: redis) if limit
+    limit = -1 if limit.to_i <= 0
+    @restrainer = Restrainer.new("DoubleRestrainer(#{name})", limit: limit, redis: redis)
     @long_running_restrainer = Restrainer.new("DoubleRestrainer(#{name}).long_running", limit: long_running_limit, redis: redis)
   end
 
   # Execute a block of code. The block will be yielded with the timeout value. If the block raises
   # a timeout error, then it will be called again with the long running timeout. The code in the block
   # must be idempotent since it can be run twice.
+  #
   # @yieldparam [Numeric] the timeout value to use in the block.
   # @raise [Restrainer::ThrottleError] if too many concurrent processes are trying to use the restraint.
-  def execute
+  def execute(&block)
     begin
-      if @restrainer
-        @restrainer.throttle do
-          yield @timeout
-        end
-      else
+      @restrainer.throttle do
         yield @timeout
       end
     rescue => e
       if @timeout_errors.any? { |error_class| e.is_a?(error_class) }
         @long_running_restrainer.throttle do
@@ -45,7 +45,47 @@
         end
       else
         raise e
       end
     end
+  end
+
+  # Get the current size of the default pool. This can be useful in
+  # collecting realtime stats about how the pool is being utilized.
+  #
+  # @return [Integer]
+  def default_pool_size
+    @restrainer.current
+  end
+
+  # Get the current size of the long running pool. This can be useful in
+  # collecting realtime stats about how the pool is being utilized.
+  #
+  # @return [Integer]
+  def long_running_pool_size
+    @long_running_restrainer.current
+  end
+
+  # Get the limit for the default pool. This will return -1 if there
+  # is no limit set on that pool.
+  #
+  # @return [Integer]
+  def default_pool_limit
+    @restrainer.limit
+  end
+
+  # Get the limit for the long running pool.
+  #
+  # @return [Integer]
+  def long_running_pool_limit
+    @long_running_restrainer.limit
+  end
+
+  # Helper method to determine if a timeout represents the long running timeout.
+  # Note that the timeout and long running timeouts need to be different values
+  # in order for this to work.
+  #
+  # @return [Boolean]
+  def long_running?(timeout)
+    timeout.to_f.round(6) == @long_running_timeout.to_f.round(6)
   end
 end