<s> """<STR_LIT>""" <EOL> import logging <EOL> import math <EOL> import redis <EOL> import time <EOL> from itertools import izip <EOL> class RateLimiter ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , conditions = None , <EOL> redis_host = '<STR_LIT:localhost>' , <EOL> redis_port = <NUM_LIT> , <EOL> redis_db = <NUM_LIT:0> , <EOL> redis_password = None , <EOL> redis_namespace = '<STR_LIT>' , <EOL> redis_url = None ) : <EOL> """<STR_LIT>""" <EOL> if redis_url : <EOL> self . redis = redis . Redis . from_url ( redis_url , db = redis_db ) <EOL> else : <EOL> self . redis = redis . Redis ( host = redis_host , port = redis_port , db = redis_db , password = redis_password ) <EOL> self . log = logging . getLogger ( '<STR_LIT>' ) <EOL> self . namespace = redis_namespace <EOL> self . conditions = [ ] <EOL> self . list_ttl = <NUM_LIT:0> <EOL> if conditions : <EOL> self . add_condition ( * conditions ) <EOL> def add_condition ( self , * conditions ) : <EOL> """<STR_LIT>""" <EOL> if len ( conditions ) == <NUM_LIT:2> and isinstance ( conditions [ <NUM_LIT:0> ] , int ) : <EOL> conditions = [ conditions ] <EOL> for condition in conditions : <EOL> if isinstance ( condition , dict ) : <EOL> requests = condition [ '<STR_LIT>' ] <EOL> seconds = condition . get ( '<STR_LIT>' , <NUM_LIT:0> ) + ( <EOL> <NUM_LIT> * ( condition . get ( '<STR_LIT>' , <NUM_LIT:0> ) + <EOL> <NUM_LIT> * ( condition . get ( '<STR_LIT>' , <NUM_LIT:0> ) + <EOL> <NUM_LIT> * condition . get ( '<STR_LIT>' , <NUM_LIT:0> ) ) ) ) <EOL> else : <EOL> requests , seconds = condition <EOL> requests = int ( requests ) <EOL> seconds = int ( seconds ) <EOL> if requests < <NUM_LIT:0> : <EOL> raise ValueError ( '<STR_LIT>' % requests ) <EOL> if seconds < <NUM_LIT:0> : <EOL> raise ValueError ( '<STR_LIT>' % seconds ) <EOL> if seconds > <NUM_LIT:0> : <EOL> if requests == <NUM_LIT:0> : <EOL> self . log . warn ( '<STR_LIT>' , requests , seconds ) <EOL> else : <EOL> self . log . debug ( '<STR_LIT>' , requests , seconds ) <EOL> self . conditions . append ( ( requests , seconds ) ) <EOL> if seconds > self . list_ttl : <EOL> self . list_ttl = seconds <EOL> else : <EOL> self . log . warn ( '<STR_LIT>' ) <EOL> self . conditions . sort ( ) <EOL> def block ( self , key , seconds = <NUM_LIT:0> , minutes = <NUM_LIT:0> , hours = <NUM_LIT:0> , days = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> seconds = seconds + <NUM_LIT> * ( minutes + <NUM_LIT> * ( hours + <NUM_LIT> * days ) ) <EOL> if not seconds : <EOL> seconds = self . list_ttl <EOL> if not seconds : <EOL> self . log . warn ( '<STR_LIT>' ) <EOL> return <NUM_LIT:0> <EOL> if not isinstance ( seconds , int ) : <EOL> seconds = int ( math . ceil ( seconds ) ) <EOL> key = '<STR_LIT::>' . join ( ( self . namespace , key , '<STR_LIT>' ) ) <EOL> self . log . warn ( '<STR_LIT>' , key , seconds ) <EOL> with self . redis . pipeline ( ) as pipe : <EOL> try : <EOL> pipe . set ( key , '<STR_LIT:1>' ) <EOL> pipe . expire ( key , seconds ) <EOL> pipe . execute ( ) <EOL> except redis . exceptions . LockError : <EOL> self . log . exception ( "<STR_LIT>" ) <EOL> return seconds <EOL> def acquire ( self , key , block = True ) : <EOL> """<STR_LIT>""" <EOL> if block : <EOL> while True : <EOL> success , wait = self . _make_ping ( key ) <EOL> if success : <EOL> return True <EOL> self . log . debug ( '<STR_LIT>' , wait ) <EOL> time . sleep ( wait ) <EOL> else : <EOL> success , wait = self . _make_ping ( key ) <EOL> return ( success , wait ) <EOL> __call__ = acquire <EOL> def _make_ping ( self , key ) : <EOL> if not self . conditions : <EOL> return True , <NUM_LIT:0.0> <EOL> min_requests , min_request_seconds = self . conditions [ <NUM_LIT:0> ] <EOL> if min_requests == <NUM_LIT:0> : <EOL> self . log . warn ( '<STR_LIT>' , key , min_requests , min_request_seconds ) <EOL> return False , min_request_seconds <EOL> log_key = '<STR_LIT::>' . join ( ( self . namespace , key , '<STR_LIT>' ) ) <EOL> block_key = '<STR_LIT::>' . join ( ( self . namespace , key , '<STR_LIT>' ) ) <EOL> lock_key = '<STR_LIT::>' . join ( ( self . namespace , key , '<STR_LIT>' ) ) <EOL> try : <EOL> with self . redis . lock ( lock_key , timeout = <NUM_LIT:10> ) : <EOL> with self . redis . pipeline ( ) as pipe : <EOL> for requests , _ in self . conditions : <EOL> pipe . lindex ( log_key , requests - <NUM_LIT:1> ) <EOL> pipe . ttl ( block_key ) <EOL> pipe . get ( block_key ) <EOL> boundry_timestamps = pipe . execute ( ) <EOL> blocked = boundry_timestamps . pop ( ) <EOL> block_ttl = boundry_timestamps . pop ( ) <EOL> if blocked is not None : <EOL> if block_ttl is None : <EOL> block_ttl = <NUM_LIT:0.5> <EOL> self . log . warn ( '<STR_LIT>' , key , block_ttl ) <EOL> return False , block_ttl <EOL> timestamp = time . time ( ) <EOL> for boundry_timestamp , ( requests , seconds ) in izip ( boundry_timestamps , self . conditions ) : <EOL> if boundry_timestamp is not None : <EOL> boundry_timestamp = float ( boundry_timestamp ) <EOL> if boundry_timestamp + seconds > timestamp : <EOL> self . log . warn ( '<STR_LIT>' , <EOL> key , requests , seconds , boundry_timestamp + seconds - timestamp ) <EOL> return False , boundry_timestamp + seconds - timestamp <EOL> with self . redis . pipeline ( ) as pipe : <EOL> pipe . lpush ( log_key , timestamp ) <EOL> max_requests , _ = self . conditions [ - <NUM_LIT:1> ] <EOL> pipe . ltrim ( log_key , <NUM_LIT:0> , max_requests - <NUM_LIT:1> ) <EOL> pipe . expire ( log_key , self . list_ttl ) <EOL> pipe . execute ( ) <EOL> except redis . exceptions . LockError : <EOL> self . log . exception ( "<STR_LIT>" ) <EOL> return True , <NUM_LIT:0.0> <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> """<STR_LIT>""" <EOL> import sys <EOL> logging . basicConfig ( format = '<STR_LIT>' , level = logging . DEBUG , stream = sys . stdout ) <EOL> log = logging . getLogger ( '<STR_LIT>' ) <EOL> key = '<STR_LIT>' <EOL> rate = RateLimiter ( conditions = ( ( <NUM_LIT:1> , <NUM_LIT:1> ) , ( <NUM_LIT:2> , <NUM_LIT:5> ) ) ) <EOL> rate . add_condition ( ( <NUM_LIT:3> , <NUM_LIT:10> ) , ( <NUM_LIT:4> , <NUM_LIT:15> ) ) <EOL> rate . add_condition ( { '<STR_LIT>' : <NUM_LIT:20> , '<STR_LIT>' : <NUM_LIT:5> } ) <EOL> rate . add_condition ( { '<STR_LIT>' : <NUM_LIT> , '<STR_LIT>' : <NUM_LIT:15> } , { '<STR_LIT>' : <NUM_LIT> , '<STR_LIT>' : <NUM_LIT:1> } ) <EOL> i = <NUM_LIT:1> <EOL> for _ in xrange ( <NUM_LIT:5> ) : <EOL> rate . acquire ( key ) <EOL> log . info ( '<STR_LIT>' , i ) <EOL> i += <NUM_LIT:1> <EOL> rate . block ( key , seconds = <NUM_LIT:10> ) <EOL> for _ in xrange ( <NUM_LIT:10> ) : <EOL> rate . acquire ( key ) <EOL> log . info ( '<STR_LIT>' , i ) <EOL> i += <NUM_LIT:1> <EOL> rate . add_condition ( <NUM_LIT:0> , <NUM_LIT:1> ) <EOL> for _ in xrange ( <NUM_LIT:5> ) : <EOL> rate ( key , block = False ) <EOL> time . sleep ( <NUM_LIT:1> ) </s>
