From 436b3c7d0c93bd6019a576425506c4c8f2963632 Mon Sep 17 00:00:00 2001 From: Mark Haines Date: Mon, 1 Sep 2014 17:54:54 +0100 Subject: [PATCH] Ratelimiter object --- synapse/api/ratelimiting.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 synapse/api/ratelimiting.py diff --git a/synapse/api/ratelimiting.py b/synapse/api/ratelimiting.py new file mode 100644 index 000000000..a730abd59 --- /dev/null +++ b/synapse/api/ratelimiting.py @@ -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