lib/simple_throttle.rb in simple_throttle-1.1.0 vs lib/simple_throttle.rb in simple_throttle-1.1.1

@@ -15,13 +15,14 @@
     local limit = tonumber(ARGV[1])
     local ttl = tonumber(ARGV[2])
     local now = ARGV[3]
     local pause_to_recover = tonumber(ARGV[4])
     local amount = tonumber(ARGV[5])
+    local cleanup = tonumber(ARGV[6])
 
     local size = redis.call('llen', list_key)
-    if size >= limit then
+    if size >= limit or (cleanup > 0 and size > 0) then
       local expired = tonumber(now) - ttl
       while size > 0 do
         local t = redis.call('lpop', list_key)
         if tonumber(t) > expired then
           redis.call('lpush', list_key, t)
@@ -150,30 +151,22 @@
   # Returns true if the limit for the throttle has not been reached yet. This method
   # will also track the throttled resource as having been invoked on each call.
   #
   # @return [Boolean]
   def allowed!
-    size = increment!
+    size = add_request(1, false)
     size <= limit
   end
 
   # Increment the throttle by the specified and return the current size. Because
   # how the throttle is implemented in Redis, the return value will always max
   # out at the throttle limit + 1 or, if the pause to recover option is set, limit + 2.
   #
   # @param amount [Integer] amount to increment the throttle by
   # @return [Integer]
   def increment!(amount = 1)
-    pause_to_recover_arg = (@pause_to_recover ? 1 : 0)
-    time_ms = (Time.now.to_f * 1000).round
-    ttl_ms = (ttl * 1000).ceil
-    self.class.send(
-      :execute_lua_script,
-      redis: redis_client,
-      keys: [redis_key],
-      args: [limit, ttl_ms, time_ms, pause_to_recover_arg, amount]
-    )
+    add_request(amount, true)
   end
 
   # Reset a throttle back to zero.
   #
   # @return [void]
@@ -216,7 +209,19 @@
     end
   end
 
   def redis_key
     "simple_throttle.#{name}"
+  end
+
+  def add_request(amount, cleanup)
+    pause_to_recover_arg = (@pause_to_recover ? 1 : 0)
+    time_ms = (Time.now.to_f * 1000).round
+    ttl_ms = (ttl * 1000).ceil
+    self.class.send(
+      :execute_lua_script,
+      redis: redis_client,
+      keys: [redis_key],
+      args: [limit, ttl_ms, time_ms, pause_to_recover_arg, amount, (cleanup ? 1 : 0)]
+    )
   end
 end