rate_limiter.rb in congestion-0.1.0

- old
+ new
@@ -82,23 +82,34 @@
       (millis / 1_000.0).ceil
     end
 
     def add_request
       unless options[:track_rejected]
-        redis.zadd key, current_time, current_time
+        add_request = redis.multi do |t|
+          t.zadd key, current_time, current_time # [0] - key added
+          t.ttl key                              # [1] - key ttl
+        end
+        # TTL is -1 if not set, https://redis.io/commands/ttl
+        if add_request[1] == -1
+          # ensure we set the expire TTL on the request key
+          # using the raw interval here after the 'get_requests' limit check
+          # should be close enough to the actual interval folks desire
+          redis.pexpire key, options[:interval]
+        end
       end
     end
 
     def get_requests
       @requests ||= redis.multi do |t|
-        t.zremrangebyscore key, 0, expired_at   # [0] - clear old requests
-        t.zcount key, '-inf', '+inf'            # [1] - number of requests
-        t.zrange key, 0, 0                      # [2] - first request
-        t.zrange key, -1, -1                    # [3] - last request
-        t.pexpire key, options[:interval]       # [4] - expire request key
-
-        # [5] - Add the request if tracking rejected
-        t.zadd(key, current_time, current_time) if options[:track_rejected]
+        t.zremrangebyscore key, 0, expired_at     # [0] - clear old requests
+        t.zcount key, '-inf', '+inf'              # [1] - number of requests
+        t.zrange key, 0, 0                        # [2] - first request
+        t.zrange key, -1, -1                      # [3] - last request
+        if options[:track_rejected]
+          t.zadd(key, current_time, current_time) # [4] - Add the request if tracking rejected
+        end
+        # ensure the TTL is set after we've added the key if tracking rejected
+        t.pexpire key, options[:interval]         # [5] - expire request key
       end
     end
   end
 end