Added Sliding window counter rate limiter class

faizanfareed · faizanfareed · commit 53ce4c949d98 · 2020-08-02T19:34:38.000+05:00
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,2 @@
+env/
+extra/
diff --git a/LICENSE b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) [2020] [Muhammad Faizan Fareed]
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/ratelimiter.py b/ratelimiter.py
@@ -0,0 +1,309 @@
+import datetime
+import time
+
+# For logs or debugging. Uncomment when required and inherit.
+# from utility import Utility
+
+
+class SlidingWindowCounterRateLimiter:
+
+    '''Sliding Window Counter Rate Limiter.'''
+
+    def __init__(self, clientid, redispipeline, rate=10, time_window_unit='minute', is_ratelimit_reset_header_allowed=True, is_2nd_RTT_allowed=False, max_no_time_window_for_deletion=5, is_log_enabled=True, per_request_increement=1):
+        '''Initilize all properties of rate limiter.'''
+
+        # redis pipeline
+        self.pipeline = redispipeline
+
+        # Increement counter by 1. Set if needed
+        self.per_request_increement = per_request_increement
+
+        self.is_ratelimit_reset_header_allowed = is_ratelimit_reset_header_allowed
+
+        # client id
+        self.clientid = clientid
+
+        # rate limit
+        self.rate = rate
+
+        # Time Window  ('per')
+        self.time_window = 1
+
+        # Unit of Time Window
+        self.time_window_unit = time_window_unit
+
+        # Total requert served
+        self.total_request_served = 0
+
+        '''
+            By default disabled.
+            Logs usage : For specific user and used  for debuging.
+            Uncomment when required.
+            
+        '''
+        # self.is_log_enabled = is_log_enabled
+
+        self.request_recieved_timestamp_in_TW_format = None
+        self.request_recieved_at = None
+
+        # For performance purposes.
+        self.is_2nd_RTT_allowed = is_2nd_RTT_allowed
+
+        # Reduce memeory footprint when old/expired fixed time window reached at some number.
+        self.max_no_time_window_for_deletion = max_no_time_window_for_deletion
+
+        # Set dynamically rate limit time window
+        self.windowTimeKwargs = {self.time_window_unit + 's': self.time_window}
+
+        # Set time window time format
+        self.__setFixedTimeWindowFormat()
+
+    def __setFixedTimeWindowFormat(self):
+        '''Set Fixed Time Window time in different format as a string.
+
+            Set Fixed Time Window format for insertion as a string.
+            Set Fixed Time Window format for ratelimit reset a as string.
+            Set expiration time for user keys.
+        '''
+
+        self.fixed_time_window_format_for_insertion = ''
+
+        if self.time_window_unit == 'second':
+
+            self.expiration_time_of_client_keys = self.time_window + self.time_window
+
+            self.fixed_time_window_format_for_insertion = "%Y:%m:%d:%H:%M:%S"
+            self.fixed_time_window_format_for_ratelimit_reset = "%Y-%m-%d-%H-%M-%S"
+
+        elif self.time_window_unit == 'minute':
+
+            self.expiration_time_of_client_keys = self.time_window * \
+                60 + (self.time_window)
+
+            self.fixed_time_window_format_for_insertion = "%Y:%m:%d:%H:%M:%S"
+            self.fixed_time_window_format_for_ratelimit_reset = "%Y-%m-%d-%H-%M"
+
+        elif self.time_window_unit == 'hour':
+
+            self.expiration_time_of_client_keys = self.time_window * \
+                3600 + (self.time_window)
+            self.fixed_time_window_format_for_insertion = "%Y:%m:%d:%H:%M"
+            self.fixed_time_window_format_for_ratelimit_reset = "%Y-%m-%d-%H"
+
+    def __setRequest_Recieved_Timestamp(self):
+        '''Set request recieved time in different formats.
+
+            Set request recieved at datetime format.
+            Set request recieved timestamp for rate limit reset time in rate limit unit format.
+            Set request recieved timestamp for increement the Fixed Time Window counter in ftW insertion format.
+        '''
+
+        self.request_recieved_at = datetime.datetime.now()
+
+        self.request_recieved_timestamp_for_ratelimit_reset = self.request_recieved_at.strftime(
+            self.fixed_time_window_format_for_ratelimit_reset)
+
+        # Get current time based on window time interval format.
+        currentTime = self.request_recieved_at.strftime(
+            self.fixed_time_window_format_for_insertion)
+        self.request_recieved_timestamp_in_TW_format = int(datetime.datetime.timestamp(datetime.datetime.strptime(
+            currentTime, self.fixed_time_window_format_for_insertion)))  # Converting into unix timestamp
+
+    def __setSlidingTimeWindowTimestamp(self):
+        '''Set sliding time window time in timestamp.
+        Subtracting rate limit time window from  request recived time, foramt in FTW_format_for_insertion format and then into timestamp.
+        '''
+
+        self.slidingTimeWindow = self.request_recieved_at - \
+            datetime.timedelta(**self.windowTimeKwargs)
+        self.slidingTimeWindow = datetime.datetime.strftime(
+            self.slidingTimeWindow, self.fixed_time_window_format_for_insertion)
+        self.slidingTimeWindow = datetime.datetime.strptime(
+            self.slidingTimeWindow, self.fixed_time_window_format_for_insertion)
+        self.slidingTimeWindow = int(
+            datetime.datetime.timestamp(self.slidingTimeWindow))
+
+    def isRequestAllowed(self):
+        '''Returns True if request allowed otherwise False.'''
+
+        # Set current  timestamp in time window   format for increement counter
+        self.__setRequest_Recieved_Timestamp()
+
+        # Increement counter into current  time window
+        self.pipeline.execute_command(
+            'HINCRBY', self.clientid, self.request_recieved_timestamp_in_TW_format, self.per_request_increement)
+
+        # Get all windows
+        self.pipeline.execute_command('HGETALL', self.clientid)
+
+        # Set expiration time
+        self.pipeline.execute_command(
+            'EXPIRE', self.clientid, self.expiration_time_of_client_keys)
+
+        if self.is_ratelimit_reset_header_allowed:
+
+            # Set each fixed time window timestamp when created first time.
+            self.pipeline.execute_command('SET', self.clientid+'-'+str(self.request_recieved_timestamp_for_ratelimit_reset),
+                                          self.__get_FTW_created_at_deltatime(), 'ex', self.expiration_time_of_client_keys, 'nx')
+            # Get back fixed time window time
+            self.pipeline.execute_command(
+                'GET', self.clientid+'-'+str(self.request_recieved_timestamp_for_ratelimit_reset))
+
+        '''
+            By default disabled.
+            If log enabled then insert each request into log.
+            Logs usage : For specific user and used  for debuging.
+            Uncomment when required.
+            
+        '''
+        # self.insertLogrequest()
+
+        # Execute Redus pipeline
+        result = self.pipeline.execute()
+
+        # Get all time windows from result list
+        time_window_list = result[1]
+
+        # Set sliding time window
+        self.__setSlidingTimeWindowTimestamp()
+
+        # Trim and count requests served in rate limit time window
+        self.__trim_old_and_count_requests_in_time_window(time_window_list)
+
+        if self.is_ratelimit_reset_header_allowed:
+
+            current_fixed_time_window_created_at = result[4]
+            self.__set_ratelimit_reset_time(
+                current_fixed_time_window_created_at)
+
+        '''
+            By default disabled.
+            Print each logs.
+            Logs for specific user and used  for debuging.
+            Uncomment when required.
+            
+        '''
+        # self.print_logs(result)
+
+        # Delete expired FTW
+        if self.__is_expired_time_windows_exists and (self.max_no_time_window_for_deletion > self.total_expired_time_windows or self.is_2nd_RTT_allowed):
+            self.pipeline.execute()
+
+        # If request in window greater then rate then return False otherwise True
+        if self.total_request_served > self.rate:
+
+            self.is_request_allowed = False
+            return False
+        else:
+            self.is_request_allowed = True
+            return True
+
+    def __get_FTW_created_at_deltatime(self):
+        '''Returns FTW created at time with addition of rate limit time window for rate limit reset header.
+        Instead of adding rate limit time window on each request i stored when FTW created first time.'''
+
+        resetTime = self.request_recieved_at + \
+            datetime.timedelta(**self.windowTimeKwargs)
+        return datetime.datetime.timestamp(resetTime)
+
+    def __set_ratelimit_reset_time(self, current_fixed_time_window_created_at):
+        '''Set ratelimit reset time. 
+        Fixed Time Window first time created time passed as a parameter. Subtracing the request recived time from FTW time. 
+        '''
+
+        self.x_ratelimit_reset = None
+        current_fixed_time_window_created_at = current_fixed_time_window_created_at.decode()
+        self.x_ratelimit_reset = datetime.datetime.fromtimestamp(
+            float(current_fixed_time_window_created_at)) - self.request_recieved_at
+
+    def get_x_ratelimit_reset(self, in_seconds=True):
+        '''Returns x-ratelimit-reset in seconds or in str format.'''
+
+        if in_seconds:
+            return self.x_ratelimit_reset.total_seconds()
+        else:
+            return self.x_ratelimit_reset
+
+    def __trim_old_and_count_requests_in_time_window(self, time_window_list):
+        '''Counts no of reqeuest served in rate limit time window, insert time windows into pipeline which are less then
+        sliding time window, set True if any expired  fixed time window exists and count total expired windows.
+        '''
+
+        self.__is_expired_time_windows_exists = False
+        self.total_expired_time_windows = 0
+
+        for window, counts in time_window_list.items():
+
+            if int(window.decode()) < self.slidingTimeWindow:  # Check each window expired or not
+
+                # Insert all expired windows into pipeline
+                self.pipeline.execute_command(
+                    'HDEL', self.clientid, int(window.decode()))
+
+                # Set if some windows expired for delete expired windows
+                self.__is_expired_time_windows_exists = True
+
+                # Count no of expired windows
+                self.total_expired_time_windows = self.total_expired_time_windows + 1
+
+            else:  # if windows not expired then count all requests
+
+                # Counts request in window size
+                self.total_request_served = self.total_request_served + \
+                    int(counts.decode())
+
+    def get_retry_after(self):
+        ''' Returns retry after header.'''
+
+        retry_after = self.request_recieved_at + \
+            datetime.timedelta(**self.windowTimeKwargs)
+        return retry_after.strftime("%Y:%m:%d:%H:%M:%S:%M")
+
+    def get_x_ratelimit_remaining(self):
+        '''Returns rate limit remaining requests.
+        Subtracting total request served from max allowed requests. 
+        '''
+
+        remaining_request = self.getMaxRequestsAllowed() - self.total_request_served
+
+        if remaining_request <= 0:
+            return 0
+        else:
+            return remaining_request
+
+    def getMaxRequestsAllowed(self):
+        '''Returns max no of requests allowed.'''
+
+        return self.rate
+
+    def getTotalRequestServedInSlidingWindow(self):
+        '''Returns total request served in rate limit time window'''
+
+        return self.total_request_served
+
+    def getHttpResponseHeaders(self):
+        '''Returns Http reseponse headers.
+
+            X-RateLimit-Limit.
+            X_RateLimit_Remaining.
+            X-RateLimit-Reset.
+
+        '''
+
+        HttpResponseHeaders = {
+
+            'X-RateLimit-Limit': self.getMaxRequestsAllowed(),  # Max requests allowed
+            'X_RateLimit_Remaining': self.get_x_ratelimit_remaining(),  # Remaing allowed requests
+
+        }
+
+        # Reset time in seconds if this header allowed
+        if self.is_ratelimit_reset_header_allowed:
+            HttpResponseHeaders['X-RateLimit-Reset'] = self.get_x_ratelimit_reset(
+                in_seconds=True)
+
+        # If status code 429 then Retry after header
+        if not self.is_request_allowed:
+            HttpResponseHeaders['Retry-after'] = self.get_retry_after()
+
+        return HttpResponseHeaders
diff --git a/utility.py b/utility.py
