Merge pull request #56 from zendesk/fvilela/after_evaluated_callback

[RED-2096] Add after_evaluated callback option

Author: fbvilela

README.md

@@ -31,6 +31,17 @@ Prop.before_throttle do |handle, key, threshold, interval|
 end
 ```
 
+## Setting an After Evaluated Callback
+
+You can define an optional callback that is invoked when a rate limit is checked. The callback will be invoked regardless
+of the result of the evaluation.
+
+```ruby
+Prop.after_evaluated do |handle, counter, options|
+  Rails.logger.info "Prop #{handle} has just been check. current value: #{counter}"
+end
+````
+
 ## Defining thresholds
 
 Example: Limit on accepted emails per hour from a given user, by defining a threshold and interval (in seconds):

lib/prop.rb

@@ -9,6 +9,6 @@ module Prop
   class << self
     extend Forwardable
     def_delegators :"Prop::Limiter", :read, :write, :cache, :cache=, :configure, :configurations, :disabled, :before_throttle
-    def_delegators :"Prop::Limiter", :throttle, :throttle!, :throttled?, :count, :query, :reset
+    def_delegators :"Prop::Limiter", :throttle, :throttle!, :throttled?, :count, :query, :reset, :after_evaluated
   end
 end

lib/prop/limiter.rb

@@ -9,7 +9,7 @@ module Prop
   class Limiter
 
     class << self
-      attr_accessor :handles, :before_throttle_callback, :cache
+      attr_accessor :handles, :before_throttle_callback, :cache, :after_evaluated_callback
 
       def read(&blk)
         raise "Use .cache = "
@@ -39,6 +39,10 @@ def before_throttle(&blk)
         self.before_throttle_callback = blk
       end
 
+      def after_evaluated(&blk)
+        self.after_evaluated_callback = blk
+      end
+
       # Public: Registers a handle for rate limiting
       #
       # handle   - the name of the handle you wish to use in your code, e.g. :login_attempt
@@ -150,13 +154,18 @@ def _throttle(strategy, handle, key, cache_key, options)
         return [false, strategy.zero_counter] if disabled?
 
         if leaky_bucket_strategy?(strategy)
-          return Prop::LeakyBucketStrategy._throttle_leaky_bucket(handle, key, cache_key, options)
+          is_over_limit, bucket_info_hash = Prop::LeakyBucketStrategy._throttle_leaky_bucket(handle, key, cache_key, options)
+          bucket_counter = bucket_info_hash.fetch(:bucket)
+          after_evaluated_callback.call(handle, bucket_counter, options.merge(bucket_info_hash)) if after_evaluated_callback
+          return [is_over_limit, bucket_info_hash]
         end
 
         counter = options.key?(:decrement) ?
           strategy.decrement(cache_key, options.fetch(:decrement), options) :
           strategy.increment(cache_key, options.fetch(:increment, 1), options)
 
+        after_evaluated_callback.call(handle, counter, options) if after_evaluated_callback
+
         if strategy.compare_threshold?(counter, :>, options)
           before_throttle_callback &&
             before_throttle_callback.call(handle, key, options[:threshold], options[:interval])

test/test_limiter.rb

@@ -67,6 +67,22 @@ def self.it_handles_concurrency(method)
         Prop.count(:something).must_equal -4
       end
 
+      describe "when there is a after_evaluated callback" do
+        it "invokes the callback" do
+          Prop.after_evaluated do |handle, counter, options|
+            @handle   = handle
+            @counter  = counter
+            @options  = options
+          end
+
+          Prop.throttle(:something)
+
+          @handle.must_equal :something
+          @counter.must_equal 1
+          @options.must_equal({ threshold: 10, interval: 10, key: "", strategy: Prop::IntervalStrategy })
+        end
+      end
+
       describe "and the threshold has been reached" do
         before { Prop::IntervalStrategy.stubs(:compare_threshold?).returns(true) }
 
@@ -190,6 +206,22 @@ def self.it_handles_concurrency(method)
           )
         end
       end
+
+      describe "when there is a after_evaluated callback" do
+        it "invokes the callback" do
+          Prop.after_evaluated do |handle, counter, options|
+            @handle   = handle
+            @counter  = counter
+            @options  = options
+          end
+
+          Prop.throttle(:something)
+
+          @handle.must_equal :something
+          @counter.must_equal 1
+          @options.keys.sort.must_equal [:bucket, :burst_rate, :interval, :key, :last_leak_time, :over_limit, :strategy, :threshold]
+        end
+      end
     end
 
     describe "#throttle!" do