<s> import time <EOL> import errno <EOL> import socket <EOL> from weakref import ref as weakref <EOL> from itertools import izip <EOL> from redis import StrictRedis <EOL> from redis . client import list_or_args <EOL> from redis . exceptions import ConnectionError <EOL> try : <EOL> from redis . exceptions import TimeoutError <EOL> except ImportError : <EOL> TimeoutError = ConnectionError <EOL> from rb . promise import Promise <EOL> from rb . poll import poll <EOL> AUTO_BATCH_COMMANDS = { <EOL> '<STR_LIT:GET>' : ( '<STR_LIT>' , True ) , <EOL> '<STR_LIT>' : ( '<STR_LIT>' , False ) , <EOL> } <EOL> def assert_open ( client ) : <EOL> if client . closed : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> def send_buffer ( buf , sock , host_id ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> timeout = sock . gettimeout ( ) <EOL> sock . setblocking ( False ) <EOL> try : <EOL> for idx , item in enumerate ( buf ) : <EOL> sent = <NUM_LIT:0> <EOL> while <NUM_LIT:1> : <EOL> try : <EOL> sent = sock . send ( item ) <EOL> except IOError as e : <EOL> if e . errno == errno . EAGAIN : <EOL> continue <EOL> elif e . errno == errno . EWOULDBLOCK : <EOL> break <EOL> raise <EOL> break <EOL> if sent < len ( item ) : <EOL> buf [ : idx + <NUM_LIT:1> ] = [ item [ sent : ] ] <EOL> break <EOL> else : <EOL> del buf [ : ] <EOL> finally : <EOL> sock . settimeout ( timeout ) <EOL> except IOError as e : <EOL> if isinstance ( e , socket . timeout ) : <EOL> raise TimeoutError ( '<STR_LIT>' <EOL> % host_id ) <EOL> raise ConnectionError ( '<STR_LIT>' <EOL> % ( host_id , e ) ) <EOL> def merge_batch ( command_name , arg_promise_tuples ) : <EOL> batch_command , list_response = AUTO_BATCH_COMMANDS [ command_name ] <EOL> if len ( arg_promise_tuples ) == <NUM_LIT:1> : <EOL> args , promise = arg_promise_tuples [ <NUM_LIT:0> ] <EOL> return command_name , args , promise <EOL> promise = Promise ( ) <EOL> @ promise . done <EOL> def on_success ( value ) : <EOL> if list_response : <EOL> for item , ( _ , promise ) in izip ( value , arg_promise_tuples ) : <EOL> promise . resolve ( item ) <EOL> else : <EOL> for _ , promise in arg_promise_tuples : <EOL> promise . resolve ( value ) <EOL> args = [ ] <EOL> for individual_args , _ in arg_promise_tuples : <EOL> args . extend ( individual_args ) <EOL> return batch_command , args , promise <EOL> def auto_batch_commands ( commands ) : <EOL> """<STR_LIT>""" <EOL> pending_batch = None <EOL> for command_name , args , promise in commands : <EOL> if command_name not in AUTO_BATCH_COMMANDS : <EOL> if pending_batch : <EOL> yield merge_batch ( * pending_batch ) <EOL> pending_batch = None <EOL> yield command_name , args , promise <EOL> continue <EOL> if pending_batch and pending_batch [ <NUM_LIT:0> ] == command_name : <EOL> pending_batch [ <NUM_LIT:1> ] . append ( ( args , promise ) ) <EOL> else : <EOL> if pending_batch : <EOL> yield merge_batch ( * pending_batch ) <EOL> pending_batch = ( command_name , [ ( args , promise ) ] ) <EOL> if pending_batch : <EOL> yield merge_batch ( * pending_batch ) <EOL> class CommandBuffer ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , host_id , connection , auto_batch = True ) : <EOL> self . host_id = host_id <EOL> self . connection = connection <EOL> self . commands = [ ] <EOL> self . pending_responses = [ ] <EOL> self . auto_batch = auto_batch <EOL> self . _send_buf = [ ] <EOL> connection . connect ( ) <EOL> @ property <EOL> def closed ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . connection is None or self . connection . _sock is None <EOL> def fileno ( self ) : <EOL> """<STR_LIT>""" <EOL> assert_open ( self ) <EOL> return self . connection . _sock . fileno ( ) <EOL> def enqueue_command ( self , command_name , args ) : <EOL> """<STR_LIT>""" <EOL> assert_open ( self ) <EOL> promise = Promise ( ) <EOL> self . commands . append ( ( command_name , args , promise ) ) <EOL> return promise <EOL> @ property <EOL> def has_pending_requests ( self ) : <EOL> """<STR_LIT>""" <EOL> return bool ( self . _send_buf or self . commands ) <EOL> def send_pending_requests ( self ) : <EOL> """<STR_LIT>""" <EOL> assert_open ( self ) <EOL> unsent_commands = self . commands <EOL> if unsent_commands : <EOL> self . commands = [ ] <EOL> if self . auto_batch : <EOL> unsent_commands = auto_batch_commands ( unsent_commands ) <EOL> buf = [ ] <EOL> for command_name , args , promise in unsent_commands : <EOL> buf . append ( ( command_name , ) + tuple ( args ) ) <EOL> self . pending_responses . append ( ( command_name , promise ) ) <EOL> cmds = self . connection . pack_commands ( buf ) <EOL> self . _send_buf . extend ( cmds ) <EOL> if not self . _send_buf : <EOL> return True <EOL> try : <EOL> send_buffer ( self . _send_buf , self . connection . _sock , self . host_id ) <EOL> except Exception : <EOL> self . connection . disconnect ( ) <EOL> raise <EOL> return not self . _send_buf <EOL> def wait_for_responses ( self , client ) : <EOL> """<STR_LIT>""" <EOL> assert_open ( self ) <EOL> if self . has_pending_requests : <EOL> raise RuntimeError ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> pending = self . pending_responses <EOL> self . pending_responses = [ ] <EOL> for command_name , promise in pending : <EOL> value = client . parse_response ( <EOL> self . connection , command_name ) <EOL> promise . resolve ( value ) <EOL> class RoutingPool ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , cluster ) : <EOL> self . cluster = cluster <EOL> def get_connection ( self , command_name , shard_hint = None ) : <EOL> host_id = shard_hint <EOL> if host_id is None : <EOL> raise RuntimeError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> real_pool = self . cluster . get_pool_for_host ( host_id ) <EOL> con = real_pool . get_connection ( command_name ) <EOL> con . __creating_pool = weakref ( real_pool ) <EOL> return con <EOL> def release ( self , connection ) : <EOL> try : <EOL> real_pool = connection . __creating_pool ( ) <EOL> except ( AttributeError , TypeError ) : <EOL> real_pool = None <EOL> if real_pool is not None : <EOL> real_pool . release ( connection ) <EOL> def disconnect ( self ) : <EOL> self . cluster . disconnect_pools ( ) <EOL> def reset ( self ) : <EOL> pass <EOL> class BaseClient ( StrictRedis ) : <EOL> pass <EOL> class RoutingBaseClient ( BaseClient ) : <EOL> def __init__ ( self , connection_pool , auto_batch = True ) : <EOL> BaseClient . __init__ ( self , connection_pool = connection_pool ) <EOL> self . auto_batch = auto_batch <EOL> def pubsub ( self , ** kwargs ) : <EOL> raise NotImplementedError ( '<STR_LIT>' ) <EOL> def pipeline ( self , transaction = True , shard_hint = None ) : <EOL> raise NotImplementedError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> def lock ( self , * args , ** kwargs ) : <EOL> raise NotImplementedError ( '<STR_LIT>' ) <EOL> class MappingClient ( RoutingBaseClient ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , connection_pool , max_concurrency = None , <EOL> auto_batch = True ) : <EOL> RoutingBaseClient . __init__ ( self , connection_pool = connection_pool , <EOL> auto_batch = auto_batch ) <EOL> self . _max_concurrency = max_concurrency <EOL> self . _cb_poll = poll ( ) <EOL> def mget ( self , keys , * args ) : <EOL> args = list_or_args ( keys , args ) <EOL> return Promise . all ( [ self . get ( arg ) for arg in args ] ) <EOL> def mset ( self , * args , ** kwargs ) : <EOL> return Promise . all ( [ self . set ( k , v ) for k , v in dict ( * args , ** kwargs ) <EOL> . iteritems ( ) ] ) . then ( lambda x : None ) <EOL> def execute_command ( self , * args ) : <EOL> router = self . connection_pool . cluster . get_router ( ) <EOL> host_id = router . get_host_for_command ( args [ <NUM_LIT:0> ] , args [ <NUM_LIT:1> : ] ) <EOL> buf = self . _get_command_buffer ( host_id , args [ <NUM_LIT:0> ] ) <EOL> return buf . enqueue_command ( args [ <NUM_LIT:0> ] , args [ <NUM_LIT:1> : ] ) <EOL> def _get_command_buffer ( self , host_id , command_name ) : <EOL> """<STR_LIT>""" <EOL> buf = self . _cb_poll . get ( host_id ) <EOL> if buf is not None : <EOL> return buf <EOL> if self . _max_concurrency is not None : <EOL> while len ( self . _cb_poll ) >= self . _max_concurrency : <EOL> self . join ( timeout = <NUM_LIT:1.0> ) <EOL> connection = self . connection_pool . get_connection ( <EOL> command_name , shard_hint = host_id ) <EOL> buf = CommandBuffer ( host_id , connection , self . auto_batch ) <EOL> self . _cb_poll . register ( host_id , buf ) <EOL> return buf <EOL> def _release_command_buffer ( self , command_buffer ) : <EOL> """<STR_LIT>""" <EOL> if command_buffer . closed : <EOL> return <EOL> self . _cb_poll . unregister ( command_buffer . host_id ) <EOL> self . connection_pool . release ( command_buffer . connection ) <EOL> command_buffer . connection = None <EOL> def join ( self , timeout = None ) : <EOL> """<STR_LIT>""" <EOL> remaining = timeout <EOL> while self . _cb_poll and ( remaining is None or remaining > <NUM_LIT:0> ) : <EOL> now = time . time ( ) <EOL> rv = self . _cb_poll . poll ( remaining ) <EOL> if remaining is not None : <EOL> remaining -= ( time . time ( ) - now ) <EOL> for command_buffer , event in rv : <EOL> if command_buffer . has_pending_requests : <EOL> if event == '<STR_LIT>' : <EOL> raise ConnectionError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> if event == '<STR_LIT>' : <EOL> command_buffer . send_pending_requests ( ) <EOL> elif event in ( '<STR_LIT>' , '<STR_LIT>' ) : <EOL> command_buffer . wait_for_responses ( self ) <EOL> self . _release_command_buffer ( command_buffer ) <EOL> if self . _cb_poll and timeout is not None : <EOL> raise TimeoutError ( '<STR_LIT>' ) <EOL> def cancel ( self ) : <EOL> """<STR_LIT>""" <EOL> for command_buffer in self . _cb_poll : <EOL> self . _release_command_buffer ( command_buffer ) <EOL> class FanoutClient ( MappingClient ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , hosts , connection_pool , max_concurrency = None , <EOL> auto_batch = True ) : <EOL> MappingClient . __init__ ( self , connection_pool , max_concurrency , <EOL> auto_batch = auto_batch ) <EOL> self . _target_hosts = hosts <EOL> self . __is_retargeted = False <EOL> self . __resolve_singular_result = False <EOL> def target ( self , hosts ) : <EOL> """<STR_LIT>""" <EOL> if self . __is_retargeted : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> rv = FanoutClient ( hosts , connection_pool = self . connection_pool , <EOL> max_concurrency = self . _max_concurrency ) <EOL> rv . _cb_poll = self . _cb_poll <EOL> rv . __is_retargeted = True <EOL> return rv <EOL> def target_key ( self , key ) : <EOL> """<STR_LIT>""" <EOL> router = self . connection_pool . cluster . get_router ( ) <EOL> host_id = router . get_host_for_key ( key ) <EOL> rv = self . target ( [ host_id ] ) <EOL> rv . __resolve_singular_result = True <EOL> return rv <EOL> def execute_command ( self , * args ) : <EOL> promises = { } <EOL> hosts = self . _target_hosts <EOL> if hosts == '<STR_LIT:all>' : <EOL> hosts = self . connection_pool . cluster . hosts . keys ( ) <EOL> elif hosts is None : <EOL> raise RuntimeError ( '<STR_LIT>' ) <EOL> for host_id in hosts : <EOL> buf = self . _get_command_buffer ( host_id , args [ <NUM_LIT:0> ] ) <EOL> promise = buf . enqueue_command ( args [ <NUM_LIT:0> ] , args [ <NUM_LIT:1> : ] ) <EOL> if self . __resolve_singular_result and len ( hosts ) == <NUM_LIT:1> : <EOL> return promise <EOL> promises [ host_id ] = promise <EOL> return Promise . all ( promises ) <EOL> class RoutingClient ( RoutingBaseClient ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , cluster , auto_batch = True ) : <EOL> RoutingBaseClient . __init__ ( self , connection_pool = RoutingPool ( cluster ) , <EOL> auto_batch = auto_batch ) <EOL> def execute_command ( self , * args , ** options ) : <EOL> pool = self . connection_pool <EOL> command_name = args [ <NUM_LIT:0> ] <EOL> command_args = args [ <NUM_LIT:1> : ] <EOL> router = self . connection_pool . cluster . get_router ( ) <EOL> host_id = router . get_host_for_command ( command_name , command_args ) <EOL> connection = pool . get_connection ( command_name , shard_hint = host_id ) <EOL> try : <EOL> connection . send_command ( * args ) <EOL> return self . parse_response ( connection , command_name , ** options ) <EOL> except ( ConnectionError , TimeoutError ) as e : <EOL> connection . disconnect ( ) <EOL> if not connection . retry_on_timeout and isinstance ( e , TimeoutError ) : <EOL> raise <EOL> connection . send_command ( * args ) <EOL> return self . parse_response ( connection , command_name , ** options ) <EOL> finally : <EOL> pool . release ( connection ) <EOL> def get_mapping_client ( self , max_concurrency = <NUM_LIT:64> , auto_batch = None ) : <EOL> """<STR_LIT>""" <EOL> if auto_batch is None : <EOL> auto_batch = self . auto_batch <EOL> return MappingClient ( connection_pool = self . connection_pool , <EOL> max_concurrency = max_concurrency , <EOL> auto_batch = auto_batch ) <EOL> def get_fanout_client ( self , hosts , max_concurrency = <NUM_LIT:64> , <EOL> auto_batch = None ) : <EOL> """<STR_LIT>""" <EOL> if auto_batch is None : <EOL> auto_batch = self . auto_batch <EOL> return FanoutClient ( hosts , connection_pool = self . connection_pool , <EOL> max_concurrency = max_concurrency , <EOL> auto_batch = auto_batch ) <EOL> def map ( self , timeout = None , max_concurrency = <NUM_LIT:64> , auto_batch = None ) : <EOL> """<STR_LIT>""" <EOL> return MapManager ( self . get_mapping_client ( max_concurrency , auto_batch ) , <EOL> timeout = timeout ) <EOL> def fanout ( self , hosts = None , timeout = None , max_concurrency = <NUM_LIT:64> , <EOL> auto_batch = None ) : <EOL> """<STR_LIT>""" <EOL> return MapManager ( self . get_fanout_client ( hosts , max_concurrency , <EOL> auto_batch ) , <EOL> timeout = timeout ) <EOL> class LocalClient ( BaseClient ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , connection_pool = None , ** kwargs ) : <EOL> if connection_pool is None : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> BaseClient . __init__ ( self , connection_pool = connection_pool , <EOL> ** kwargs ) <EOL> class MapManager ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , mapping_client , timeout ) : <EOL> self . mapping_client = mapping_client <EOL> self . timeout = timeout <EOL> self . entered = None <EOL> def __enter__ ( self ) : <EOL> self . entered = time . time ( ) <EOL> return self . mapping_client <EOL> def __exit__ ( self , exc_type , exc_value , tb ) : <EOL> if exc_type is not None : <EOL> self . mapping_client . cancel ( ) <EOL> else : <EOL> timeout = self . timeout <EOL> if timeout is not None : <EOL> timeout = max ( <NUM_LIT:1> , timeout - ( time . time ( ) - self . entered ) ) <EOL> self . mapping_client . join ( timeout = timeout ) </s>
