|
|
@@ -0,0 +1,36 @@ |
|
|
|
import collections |
|
|
|
|
|
|
|
|
|
|
|
class Ratelimiter(object): |
|
|
|
|
|
|
|
def __init__(self): |
|
|
|
self.message_counts = collections.OrderedDict() |
|
|
|
|
|
|
|
def prune_message_counts(self, time_now): |
|
|
|
for user_id in self.message_counts.keys(): |
|
|
|
message_count, time_start, msg_rate_hz = ( |
|
|
|
self.message_counts[user_id] |
|
|
|
) |
|
|
|
time_delta = time_now - time_start |
|
|
|
if message_count - time_delta * msg_rate_hz > 0: |
|
|
|
break |
|
|
|
else: |
|
|
|
del self.message_counts[user_id] |
|
|
|
|
|
|
|
def send_message(self, user_id, time_now, msg_rate_hz, burst_count): |
|
|
|
self.prune_message_counts(time_now) |
|
|
|
message_count, time_start, _ignored = self.message_counts.pop( |
|
|
|
user_id, (0., time_now, None), |
|
|
|
) |
|
|
|
time_delta = time_now - time_start |
|
|
|
if message_count - time_delta * msg_rate_hz < 0: |
|
|
|
a |
|
|
|
if message_count - (time_now - time_start) * msg_rate_hz > burst_count: |
|
|
|
allowed = False |
|
|
|
else: |
|
|
|
allowed = True |
|
|
|
message_count += 1 |
|
|
|
self.message_counts[user_id] = ( |
|
|
|
message_count, time_start, msg_rate_hz |
|
|
|
) |
|
|
|
return allowed |