├── LICENSE ├── README.md ├── ratelimit-stats.py └── ratelimit.py /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2013 Jason Foote, DomainTools LLC 2 | 3 | Permission is hereby granted, free of charge, to any person 4 | obtaining a copy of this software and associated documentation 5 | files (the "Software"), to deal in the Software without 6 | restriction, including without limitation the rights to use, 7 | copy, modify, merge, publish, distribute, sublicense, and/or sell 8 | copies of the Software, and to permit persons to whom the 9 | Software is furnished to do so, subject to the following 10 | conditions: 11 | 12 | The above copyright notice and this permission notice shall be 13 | included in all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES 17 | OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 18 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 19 | HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 20 | WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 21 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 22 | OTHER DEALINGS IN THE SOFTWARE. 23 | 24 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | rate-limit 2 | ========== 3 | 4 | Implementation of a Redis rate limiter in Python 5 | 6 | ## class __RateLimiter__ 7 | **************************************** 8 | RateLimiter is used to define one or more rate limit rules. 9 | These rules are checked on .acquire() and we either return True or False based on if we can make the request, 10 | or we can block until we make the request. 11 | Manual blocks are also supported with the block method. 12 | 13 | 14 | ### __methods__ 15 | **************************************** 16 | 17 | #### def \__init\__(self, conditions=None, redis_host='localhost', redis_port=6379, redis_db=0, redis_password=None, redis_namespace='ratelimiter'): 18 | 19 | Initalize an instance of a RateLimiter 20 | 21 | conditions - list or tuple of rate limit rules 22 | redis_host - Redis host to use 23 | redis_port - Redis port (if different than default 6379) 24 | redis_db - Redis DB to use (if different than 0) 25 | redis_password - Redis password (if needed) 26 | redis_namespace - Redis key namespace 27 | 28 | #### def __acquire__(self, key, block=True): 29 | 30 | Tests whether we can make a request, or if we are currently being limited 31 | 32 | key - key to track what to rate limit 33 | block - Whether to wait until we can make the request 34 | 35 | #### def __add_condition__(self, *conditions): 36 | 37 | Adds one or more conditions to this RateLimiter instance 38 | Conditions can be given as: 39 | 40 | add_condition(1, 10) 41 | add_condition((1, 10)) 42 | add_condition((1, 10), (30, 600)) 43 | add_condition({'requests': 1, 'seconds': 10}) 44 | add_condition({'requests': 1, 'seconds': 10}, {'requests': 200, 'hours': 6}) 45 | 46 | dict can contain 'seconds', 'minutes', 'hours', and 'days' time period parameters 47 | 48 | #### def __block__(self, key, seconds=0, minutes=0, hours=0, days=0): 49 | 50 | Set manual block for key for a period of time 51 | 52 | key - key to track what to rate limit 53 | Time parameters are added together and is the period to block for 54 | seconds 55 | minutes 56 | hours 57 | days 58 | -------------------------------------------------------------------------------- /ratelimit-stats.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2.7 2 | 3 | """ 4 | This module gathers stats on running RateLimiter instances 5 | 6 | To note. If we are not keeping keys for a whole day RateLimit entries will be cleaned up and not counted in this stats script. 7 | 8 | """ 9 | import sys 10 | import time 11 | import collections 12 | import redis 13 | import argparse 14 | import logging 15 | 16 | def ratelimiter_stats(redis_host, namespace='ratelimiter', min_daily_count=None): 17 | """ 18 | Gather and show statistics on running RateLimiter instances 19 | """ 20 | conn = redis.Redis(redis_host) 21 | now = time.time() 22 | log = logging.getLogger('stats.ratelimiter') 23 | 24 | # remove namespace leaving key 25 | l_namespace = len(namespace) 26 | 27 | all_keys = set() 28 | 29 | rl5m = collections.Counter() 30 | rl15m = collections.Counter() 31 | rl1h = collections.Counter() 32 | rl1d = collections.Counter() 33 | 34 | tot5m = 0 35 | tot15m = 0 36 | tot1h = 0 37 | tot1d = 0 38 | 39 | keys = conn.keys(namespace + '*:log') 40 | l_keys = len(keys) 41 | for i, wholekey in enumerate(keys, start=1): 42 | # remove namespace and :log 43 | key = wholekey[l_namespace:-4] 44 | 45 | # only print progress every 1000 keys 46 | if not i % 1000: 47 | log.debug('gathering stats (%d / %d)', i, l_keys) 48 | 49 | all_keys.add(key) 50 | 51 | timestamps = conn.lrange(wholekey, 0, -1) 52 | 53 | for request_time in timestamps: 54 | request_ago = now - float(request_time) 55 | if request_ago < 24 * 60 * 60.0: 56 | rl1d[key] += 1 57 | tot1d += 1 58 | if request_ago < 60 * 60.0: 59 | rl1h[key] += 1 60 | tot1h += 1 61 | if request_ago < 15 * 60.0: 62 | rl15m[key] += 1 63 | tot15m += 1 64 | if request_ago < 5 * 60.0: 65 | rl5m[key] += 1 66 | tot5m += 1 67 | 68 | 69 | # order keys by one day count 70 | stat_log = logging.getLogger('stats.ratelimiter.keys') 71 | if all_keys: 72 | all_keys = sorted(all_keys, key=lambda k: (rl1d[k], rl1h[k], rl15m[k], rl5m[k])) # sort most requested ascenting 73 | #all_keys = sorted(all_keys, key=lambda k: rl5m[k]*12*24 + rl15m[k]*4*24 + rl1h[k]*24 + rl1d[k]) # sort most requests ascenting weighted to show most recent last 74 | #all_keys = sorted(all_keys, key=lambda k: (rl5m[k], rl15m[k], rl1h[k], rl1d[k])) # sort most recent requests ascending 75 | max_key_length = max(len(key) for key in all_keys) 76 | stat_log.debug('%s\t 5m\t 15m\t 1h\t 1d', 'key'.ljust(max_key_length)) 77 | for key in all_keys: 78 | if not min_daily_count or rl1d[key] >= min_daily_count: 79 | stat_log.info('(%s)\t%5d\t%5d\t%5d\t%5d', key.ljust(max_key_length), rl5m[key], rl15m[key], rl1h[key], rl1d[key]) 80 | 81 | stat_log = logging.getLogger('stats.ratelimiter.totals') 82 | stat_log.debug(' 5m\t 15m\t 1h\t 1d') 83 | stat_log.info('%5d\t%5d\t%5d\t%5d', tot5m, tot15m, tot1h, tot1d) 84 | 85 | 86 | def manual_blocks(redis_host, namespace): 87 | """ 88 | Gather and show information on manual blocks 89 | """ 90 | # manual blocks 91 | conn = redis.Redis(redis_host) 92 | l_namespace = len(namespace) 93 | blocks = {} 94 | log = logging.getLogger('stats.blocks') 95 | keys = conn.keys(namespace + '*:block') 96 | if keys: 97 | for wholekey in keys: 98 | # remove namespace and :block 99 | ttl = conn.ttl(wholekey) 100 | # ttl will be None for last second of its life 101 | if ttl is None: 102 | ttl = 1 103 | key = wholekey[l_namespace:-6] 104 | blocks[key] = ttl 105 | 106 | blocks = sorted(blocks.iteritems(), key=lambda k: k[1]) 107 | 108 | log.debug('block_key hours:minutes:seconds remaining') 109 | for key, seconds in blocks: 110 | 111 | hours = seconds // 3600 112 | seconds = seconds % 3600 113 | minutes = seconds // 60 114 | seconds = seconds % 60 115 | 116 | log.info('%s\t\t%02d:%02d:%02d', key, hours, minutes, seconds) 117 | 118 | 119 | if __name__ == '__main__': 120 | parser = argparse.ArgumentParser(description='RateLimiter Stats') 121 | parser.add_argument('--host', action='store', dest='host', default='localhost', help='Redis Server Host (default %(default)s)') 122 | parser.add_argument('--namespace', action='store', dest='namespace', default='ratelimiter:', help='ratelimiter key namespace (default "%(default)s")') 123 | parser.add_argument('-r', '--hide-ratelimit-stats', action='store_true', dest='hide_ratelimit', default=False, help='Hide stats on RateLimiter requests (default %(default)s)') 124 | parser.add_argument('-b', '--hide-manual-blocks', action='store_true', dest='hide_manual_block', default=False, help='Hide stats on force blocked requests (default %(default)s)') 125 | parser.add_argument('-c', '--min-daily-count', action='store', dest='min_daily_count', default=False, type=int, help='Hide requests with daily count less than n (default %(default)s)') 126 | 127 | args = parser.parse_args() 128 | logging.basicConfig(format='%(asctime)s %(levelname)s %(name)s %(message)s', level=logging.DEBUG, stream=sys.stdout) 129 | 130 | log = logging.getLogger('stats') 131 | 132 | if not args.hide_ratelimit: 133 | try: 134 | ratelimiter_stats(args.host, args.namespace, args.min_daily_count) 135 | except Exception, e: 136 | log.exception(e) 137 | 138 | if not args.hide_manual_block: 139 | try: 140 | manual_blocks(args.host, args.namespace) 141 | except Exception, e: 142 | log.exception(e) 143 | 144 | -------------------------------------------------------------------------------- /ratelimit.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2.7 2 | """ 3 | This module implements a RateLimiter class. 4 | RateLimiter is a Redis backed object used to define one or more rules to rate limit requests. 5 | 6 | This module can be run to show an example of a running RateLimiter instance. 7 | """ 8 | 9 | import logging 10 | import math 11 | import redis 12 | import time 13 | from itertools import izip 14 | 15 | class RateLimiter(object): 16 | """ 17 | RateLimiter is used to define one or more rate limit rules. 18 | These rules are checked on .acquire() and we either return True or False based on if we can make the request, 19 | or we can block until we make the request. 20 | Manual blocks are also supported with the block method. 21 | """ 22 | 23 | def __init__(self, conditions=None, 24 | redis_host='localhost', 25 | redis_port=6379, 26 | redis_db=0, 27 | redis_password=None, 28 | redis_namespace='ratelimiter'): 29 | """ 30 | Initalize an instance of a RateLimiter 31 | 32 | conditions - list or tuple of rate limit rules 33 | redis_host - Redis host to use 34 | redis_port - Redis port (if different than default 6379) 35 | redis_db - Redis DB to use (if different than 0) 36 | redis_password - Redis password (if needed) 37 | redis_namespace - Redis key namespace 38 | """ 39 | 40 | self.redis = redis.Redis(host=redis_host, port=redis_port, db=redis_db, password=redis_password) 41 | self.log = logging.getLogger('RateLimiter') 42 | self.namespace = redis_namespace 43 | self.conditions = [] 44 | self.list_ttl = 0 45 | 46 | if conditions: 47 | self.add_condition(*conditions) 48 | 49 | def add_condition(self, *conditions): 50 | """ 51 | Adds one or more conditions to this RateLimiter instance 52 | Conditions can be given as: 53 | add_condition(1, 10) 54 | add_condition((1, 10)) 55 | add_condition((1, 10), (30, 600)) 56 | add_condition({'requests': 1, 'seconds': 10}) 57 | add_condition({'requests': 1, 'seconds': 10}, {'requests': 200, 'hours': 6}) 58 | 59 | dict can contain 'seconds', 'minutes', 'hours', and 'days' time period parameters 60 | """ 61 | # allow add_condition(1,2) as well as add_condition((1,2)) 62 | if len(conditions) == 2 and isinstance(conditions[0], int): 63 | conditions = [conditions] 64 | 65 | for condition in conditions: 66 | if isinstance(condition, dict): 67 | requests = condition['requests'] 68 | seconds = condition.get('seconds', 0) + ( 69 | 60 * (condition.get('minutes', 0) + 70 | 60 * (condition.get('hours', 0) + 71 | 24 * condition.get('days', 0)))) 72 | else: 73 | requests, seconds = condition 74 | 75 | # requests and seconds always a positive integer 76 | requests = int(requests) 77 | seconds = int(seconds) 78 | 79 | if requests < 0: 80 | raise ValueError('negative number of requests (%s)' % requests) 81 | if seconds < 0: 82 | raise ValueError('negative time period given (%s)' % seconds) 83 | 84 | if seconds > 0: 85 | if requests == 0: 86 | self.log.warn('added block all condition (%s/%s)', requests, seconds) 87 | else: 88 | self.log.debug('added condition (%s/%s)', requests, seconds) 89 | 90 | self.conditions.append((requests, seconds)) 91 | 92 | if seconds > self.list_ttl: 93 | self.list_ttl = seconds 94 | else: 95 | self.log.warn('time period of 0 seconds. not adding condition') 96 | 97 | # sort by requests so we query redis list in order as well as know max and min requests by position 98 | self.conditions.sort() 99 | 100 | def block(self, key, seconds=0, minutes=0, hours=0, days=0): 101 | """ 102 | Set manual block for key for a period of time 103 | key - key to track what to rate limit 104 | Time parameters are added together and is the period to block for 105 | seconds 106 | minutes 107 | hours 108 | days 109 | """ 110 | seconds = seconds + 60 * (minutes + 60 * (hours + 24 * days)) 111 | # default to largest time period we are limiting by 112 | if not seconds: 113 | seconds = self.list_ttl 114 | 115 | if not seconds: 116 | self.log.warn('block called but no default block time. not blocking') 117 | return 0 118 | 119 | if not isinstance(seconds, int): 120 | seconds = int(math.ceil(seconds)) 121 | 122 | key = ':'.join((self.namespace, key, 'block')) 123 | self.log.warn('block key (%s) for %ds', key, seconds) 124 | with self.redis.pipeline() as pipe: 125 | pipe.set(key, '1') 126 | pipe.expire(key, seconds) 127 | pipe.execute() 128 | 129 | return seconds 130 | 131 | def acquire(self, key, block=True): 132 | """ 133 | Tests whether we can make a request, or if we are currently being limited 134 | key - key to track what to rate limit 135 | block - Whether to wait until we can make the request 136 | """ 137 | if block: 138 | while True: 139 | success, wait = self._make_ping(key) 140 | if success: 141 | return True 142 | self.log.debug('blocking acquire sleeping for %.1fs', wait) 143 | time.sleep(wait) 144 | else: 145 | success, wait = self._make_ping(key) 146 | return success 147 | 148 | # alternative acquire interface ratelimiter(key) 149 | __call__ = acquire 150 | 151 | def _make_ping(self, key): 152 | 153 | # shortcut if no configured conditions 154 | if not self.conditions: 155 | return True, 0.0 156 | 157 | # short cut if we are limiting to 0 requests 158 | min_requests, min_request_seconds = self.conditions[0] 159 | if min_requests == 0: 160 | self.log.warn('(%s) hit block all limit (%s/%s)', key, min_requests, min_request_seconds) 161 | return False, min_request_seconds 162 | 163 | 164 | log_key = ':'.join((self.namespace, key, 'log')) 165 | block_key = ':'.join((self.namespace, key, 'block')) 166 | lock_key = ':'.join((self.namespace, key, 'lock')) 167 | 168 | with self.redis.lock(lock_key, timeout=10): 169 | 170 | with self.redis.pipeline() as pipe: 171 | for requests, _ in self.conditions: 172 | pipe.lindex(log_key, requests-1) # subtract 1 as 0 indexed 173 | 174 | # check manual block keys 175 | pipe.ttl(block_key) 176 | pipe.get(block_key) 177 | boundry_timestamps = pipe.execute() 178 | 179 | blocked = boundry_timestamps.pop() 180 | block_ttl = boundry_timestamps.pop() 181 | 182 | if blocked is not None: 183 | # block_ttl is None for last second of a keys life. set min of 0.5 184 | if block_ttl is None: 185 | block_ttl = 0.5 186 | self.log.warn('(%s) hit manual block. %ss remaining', key, block_ttl) 187 | return False, block_ttl 188 | 189 | timestamp = time.time() 190 | 191 | for boundry_timestamp, (requests, seconds) in izip(boundry_timestamps, self.conditions): 192 | # if we dont yet have n number of requests boundry_timestamp will be None and this condition wont be limiting 193 | if boundry_timestamp is not None: 194 | boundry_timestamp = float(boundry_timestamp) 195 | if boundry_timestamp + seconds > timestamp: 196 | self.log.warn('(%s) hit limit (%s/%s) time to allow %.1fs', 197 | key, requests, seconds, boundry_timestamp + seconds - timestamp) 198 | return False, boundry_timestamp + seconds - timestamp 199 | 200 | # record our success 201 | with self.redis.pipeline() as pipe: 202 | pipe.lpush(log_key, timestamp) 203 | max_requests, _ = self.conditions[-1] 204 | pipe.ltrim(log_key, 0, max_requests-1) # 0 indexed so subtract 1 205 | # if we never use this key again, let it fall out of the DB after max seconds has past 206 | pipe.expire(log_key, self.list_ttl) 207 | pipe.execute() 208 | 209 | return True, 0.0 210 | 211 | 212 | if __name__ == '__main__': 213 | """ 214 | This is an example of rate limiting using the RateLimiter class 215 | """ 216 | import sys 217 | logging.basicConfig(format='%(asctime)s %(process)s %(levelname)s %(name)s %(message)s', level=logging.DEBUG, stream=sys.stdout) 218 | log = logging.getLogger('ratelimit.main') 219 | key = 'TestRateLimiter' 220 | 221 | rate = RateLimiter(conditions=((1, 1), (2,5))) 222 | rate.add_condition((3, 10), (4, 15)) 223 | rate.add_condition({'requests':20, 'minutes':5}) 224 | rate.add_condition({'requests':40, 'minutes':15}, {'requests':400, 'days':1}) 225 | 226 | i = 1 227 | for _ in xrange(5): 228 | rate.acquire(key) 229 | log.info('*************** ping %d ***************', i) 230 | i+=1 231 | 232 | rate.block(key, seconds=10) 233 | for _ in xrange(10): 234 | rate.acquire(key) 235 | log.info('*************** ping %d ***************', i) 236 | i+=1 237 | 238 | # block all keys 239 | rate.add_condition(0, 1) 240 | 241 | for _ in xrange(5): 242 | rate(key, block=False) # alternative interface 243 | time.sleep(1) 244 | 245 | --------------------------------------------------------------------------------