lib/prop/limiter.rb in prop-2.0.4 vs lib/prop/limiter.rb in prop-2.1.0

@@ -68,25 +68,13 @@
       # 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
       # (optional) a block of code that this throttle is guarding
       #
       # Returns true if the threshold for this handle has been reached, else returns false
-      def throttle(handle, key = nil, options = {})
-        return false if disabled?
-
+      def throttle(handle, key = nil, options = {}, &block)
         options, cache_key = prepare(handle, key, options)
-        counter = @strategy.increment(cache_key, options)
-
-        if @strategy.compare_threshold?(counter, :>, options)
-          before_throttle_callback &&
-            before_throttle_callback.call(handle, key, options[:threshold], options[:interval])
-
-          true
-        else
-          yield if block_given?
-          false
-        end
+        _throttle(handle, key, cache_key, options, &block).first
       end
 
       # Public: Records a single action for the given handle/key combination.
       #
       # handle  - the registered handle associated with the action
@@ -94,18 +82,23 @@
       # options - request specific overrides to the defaults configured for this handle
       # (optional) a block of code that this throttle is guarding
       #
       # Raises Prop::RateLimited if the threshold for this handle has been reached
       # Returns the value of the block if given a such, otherwise the current count of the throttle
-      def throttle!(handle, key = nil, options = {})
+      def throttle!(handle, key = nil, options = {}, &block)
         options, cache_key = prepare(handle, key, options)
+        throttled, counter = _throttle(handle, key, cache_key, options, &block)
 
-        if throttle(handle, key, options)
-          raise Prop::RateLimited.new(options.merge(cache_key: cache_key, handle: handle))
+        if throttled
+          raise Prop::RateLimited.new(options.merge(
+            cache_key: cache_key,
+            handle: handle,
+            first_throttled: (throttled == :first_throttled)
+          ))
         end
 
-        block_given? ? yield : @strategy.counter(cache_key, options)
+        block_given? ? yield : counter
       end
 
       # Public: Is the given handle/key combination currently throttled ?
       #
       # handle   - the throttle identifier
@@ -145,9 +138,31 @@
         @handles ||= {}
       end
       alias :configurations :handles
 
       private
+
+      def _throttle(handle, key, cache_key, options)
+        return [false, @strategy.zero_counter] if disabled?
+
+        counter = @strategy.increment(cache_key, options)
+
+        if @strategy.compare_threshold?(counter, :>, options)
+          before_throttle_callback &&
+            before_throttle_callback.call(handle, key, options[:threshold], options[:interval])
+
+          result = if options[:first_throttled] && @strategy.first_throttled?(counter, options)
+            :first_throttled
+          else
+            true
+          end
+
+          [result, counter]
+        else
+          yield if block_given?
+          [false, counter]
+        end
+      end
 
       def disabled?
         !!@disabled
       end