limiter.rb in prop-1.2.0

- old
+ new

@@ -1,8 +1,10 @@
 require 'prop/rate_limited'
 require 'prop/key'
 require 'prop/options'
+require 'prop/interval_strategy'
+require 'prop/leaky_bucket_strategy'
 
 module Prop
   class Limiter
 
     class << self
@@ -52,22 +54,21 @@
       # (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 = {})
         options, cache_key = prepare(handle, key, options)
-        counter = reader.call(cache_key).to_i
+        counter = @strategy.counter(cache_key, options)
 
         unless disabled?
-          if at_threshold?(counter, options[:threshold])
+          if @strategy.at_threshold?(counter, options)
             unless before_throttle_callback.nil?
               before_throttle_callback.call(handle, key, options[:threshold], options[:interval])
             end
 
             true
           else
-            increment = options.key?(:increment) ? options[:increment].to_i : 1
-            writer.call(cache_key, counter + increment)
+            @strategy.increment(cache_key, options, counter)
 
             yield if block_given?
 
             false
           end
@@ -88,67 +89,67 @@
 
         if throttle(handle, key, options)
           raise Prop::RateLimited.new(options.merge(:cache_key => cache_key, :handle => handle))
         end
 
-        block_given? ? yield : reader.call(cache_key).to_i
+        block_given? ? yield : @strategy.counter(cache_key, options)
       end
 
       # Public: Allows to query whether the given handle/key combination is currently throttled
       #
       # handle   - the throttle identifier
       # key      - the associated key
       #
       # Returns true if a call to `throttle!` with same parameters would raise, otherwise false
       def throttled?(handle, key = nil, options = {})
         options, cache_key = prepare(handle, key, options)
-        reader.call(cache_key).to_i >= options[:threshold]
+        counter = @strategy.counter(cache_key, options)
+        @strategy.at_threshold?(counter, options)
       end
 
       # Public: Resets a specific throttle
       #
       # handle   - the throttle identifier
       # key      - the associated key
       #
       # Returns nothing
       def reset(handle, key = nil, options = {})
         options, cache_key = prepare(handle, key, options)
-        writer.call(cache_key, 0)
+        @strategy.reset(cache_key)
       end
 
       # Public: Counts the number of times the given handle/key combination has been hit in the current window
       #
       # handle   - the throttle identifier
       # key      - the associated key
       #
       # Returns a count of hits in the current window
       def count(handle, key = nil, options = {})
         options, cache_key = prepare(handle, key, options)
-        reader.call(cache_key).to_i
+        @strategy.counter(cache_key, options)
       end
       alias :query :count
 
       def handles
         @handles ||= {}
       end
       alias :configurations :handles
 
       private
 
-      def at_threshold?(mark, threshold)
-        mark >= threshold
-      end
-
       def disabled?
         !!@disabled
       end
 
       def prepare(handle, key, params)
         raise RuntimeError.new("No such handle configured: #{handle.inspect}") unless (handles || {}).key?(handle)
 
         defaults  = handles[handle]
         options   = Prop::Options.build(:key => key, :params => params, :defaults => defaults)
-        cache_key = Prop::Key.build(:key => key, :handle => handle, :interval => options[:interval])
+
+        @strategy = options.fetch(:strategy)
+
+        cache_key = @strategy.build(:key => key, :handle => handle, :interval => options[:interval])
 
         [ options, cache_key ]
       end
     end
   end