lib/prop/limiter.rb in prop-0.7.2 vs lib/prop/limiter.rb in prop-0.7.3

@@ -43,23 +43,30 @@
       # Public: Records a single action for the given handle/key combination.
       #
       # handle  - the registered handle associated with the action
       # key     - a custom request specific key, e.g. [ account.id, "download", request.remote_ip ]
       # options - request specific overrides to the defaults configured for this handle
+      # blk     - an optional block of code that this throttle is guarding
       #
       # Raises Prop::RateLimited if the number if the threshold for this handle has been reached
-      def throttle!(handle, key = nil, options = {})
+      # Returns the value of the block if given a such, otherwise the current count of the throttle
+      def throttle!(handle, key = nil, options = {}, &blk)
         options, cache_key = prepare(handle, key, options)
-
         counter = reader.call(cache_key).to_i
 
-        return counter if disabled?
+        unless disabled?
+          if at_threshold?(counter, options[:threshold])
+            raise Prop::RateLimited.new(options.merge(:cache_key => cache_key, :handle => handle))
+          else
+            counter = writer.call(cache_key, counter + [ 1, options[:increment].to_i ].max)
+          end
+        end
 
-        if counter >= options[:threshold]
-          raise Prop::RateLimited.new(options.merge(:cache_key => cache_key, :handle => handle))
+        if block_given?
+          yield
         else
-          writer.call(cache_key, counter + [ 1, options[:increment].to_i ].max)
+          counter
         end
       end
 
       # Public: Allows to query whether the given handle/key combination is currently throttled
       #
@@ -94,9 +101,13 @@
         reader.call(cache_key).to_i
       end
       alias :query :count
 
       private
+
+      def at_threshold?(mark, threshold)
+        mark >= threshold
+      end
 
       def disabled?
         !!@disabled
       end