Simple rate limiting class in Ruby

Question

The idea of this class is I would be able to limit the use of an API/method by sending it a key which would be an IP address or username and it would tell me if they are under the rate limit.

require 'time'
class RateLimit
 def initialize(time, limit)
 @reset_on = Time.now + time
 @limit = limit
 @keys = {}
 end
 def increment(key)
 reset if Time.now >= @reset_on
 if @keys[key] == nil
 @keys[key] = 1
 else
 @keys[key] += 1
 end
 return @keys[key] >= @limit
 end
 def current_level(key)
 reset if Time.now >= @reset_on
 return @keys[key]
 end
 def is_limited?(key)
 reset if Time.now >= @reset_on
 return @keys[key] != nil && @keys[key] >= @limit
 end
 def reset
 @reset_on = Time.now + @limit
 @keys = {}
 end
end

To me this seems like some pretty clean code but are there any issues I have missed or things that could be done better?

• Flambino: This seems, well, broken. You check for individual keys, but you reset the entire keys hash. You treat @limit as a simple numerical limit, but also as a time offset? And you only set the reset_on time when resetting, so if key A is added first, and key B added later, then when things are reset, both keys are removed, even though B hasn't been there as long as A....
• Qwertie: Ah yes there is a bug, the first line of the reset method should be adding the time variable instead of limit. I was aware that all the keys get reset at the same time but I didn't have a good solution to fixing it. Now I think about it I could store the time of the first use with the key and check that that to see if just that key should be reset.
• Flambino: That'd be one way, but that - by itself - only gets you time-based throttling, not number-of-attempts-based. If you want to throttle based on number of attempts over a given period of time, you need to store attempts and their time. If you just store the count, you lose the time information, and vice-versa. So you need to store both somehow and a way to analyse those data. Until you have that, however, the code is technically broken, and not quite ready for a review

Answer 1

So there is a flaw in this approach. It limits requests based on the current time, and not by how many requests happened in the last amount of time. Let me give you an example. We'll start at time 0 with a limit of 10 requests per every seconds.

Time | Requests (total)
--------------
0 | 0
1 | 0
2 | 0
3 | 0
4 | 0
5 | 0
6 | 0
7 | 1
8 | 3
9 | 10 # requests start getting limited
10 | 10
# your request tracking resets here and requests are no longer limited
11 | 15
12 | 18
13 | 20
14 | 20
...

Do you see the issue? Between seconds 9 to 13 (4 seconds), we had a total of 17 requests when only want to allow 10 requests per every 10 seconds.

This shows that the initial approach is flawed.

In terms of code quality, I see a few things that could be improved:

1) Variable and method names:

@keys = {} is not immediately clear what it holds. However, if it were called requests_per_ip_address I would know. Long variable names better if they describe what they're holding.

is_limited? could just be limited?

2) Working with nil you should use #nil?

For example: @keys[key] == nil could be @keys.nil? and @keys[key] != nil could be !@keys[key].nil?

3) Don't return at the end of methods. The last line or branch of a method is what gets returned. The explicit return is not necessary.

4) Use private. Anyone can call any method on your rate limiter, but we only want a couple of methods to be public.

5) Use a linter. Check out RuboCop. You will be surprised by all of the style rules you break!

Here is an example of how I might write this class:

class RateLimit
 def initialize(maximum_requests, time_limit)
 @maximum_requests = maximum_requests
 @time_limit = time_limit
 @user_requests = Hash.new { |h, k| h[k] = [] }
 end
 def increment(ip_address)
 return false unless user_under_limit? ip_address
 @user_requests[ip_address] << Time.now
 true
 end
 def user_under_limit?(ip_address)
 remove_first_expired_request ip_address
 @user_requests[ip_address].count < @maximum_requests
 end
 private
 def remove_first_expired_request(ip_address)
 # We can just remove the first request if it is expired
 # and that will make room for a new request
 if @user_requests[ip_address].first < Time.now - @time_limit }
 @user_requests[ip_address].shift 
 end
 end
end

• Qwertie: Thanks for the review! The method of handling removing expired requests is quite nice. I intentionally left keys fairly vague because keys really could be almost anything. In the project I was using it in the key was a group ID but it could be an ip address or any other identifier.