<s> """<STR_LIT>""" <EOL> import datetime <EOL> import functools <EOL> from inspect import getmembers <EOL> from inspect import ismethod <EOL> import logging <EOL> import os <EOL> import random <EOL> import socket <EOL> import sys <EOL> import time <EOL> from thrift . protocol . TBinaryProtocol import TBinaryProtocolAcceleratedFactory <EOL> from thrift . Thrift import TApplicationException <EOL> from thrift . Thrift import TException <EOL> from thrift . transport . TSocket import TTransportException <EOL> from thrift . transport . TSSLSocket import TSSLSocket <EOL> from thrift . transport . TTransport import TBufferedTransport <EOL> from . . kazoo_utils . hosts import RandomHostSelector <EOL> from . base_thrift_exceptions import ThriftConnectionTimeoutError <EOL> from . connection_pool import ConnectionPool <EOL> from . retry_policy import RetryPolicy <EOL> from . TNoDelaySocket import TNoDelaySocket <EOL> DEFAULT_RETRY_POLICY = RetryPolicy ( ) <EOL> MILLIS_PER_SEC = <NUM_LIT> <EOL> log = logging . getLogger ( __name__ ) <EOL> NUM_REQUESTS_FOR_CONNECTION_TEARDOWN = os . environ . get ( <EOL> '<STR_LIT>' , <NUM_LIT:50> ) <EOL> SECS_FOR_CONNECTION_TEARDOWN = os . environ . get ( <EOL> '<STR_LIT>' , <NUM_LIT:10> ) <EOL> DEFAULT_THRIFT_PROTOCOL_FACTORY = TBinaryProtocolAcceleratedFactory ( ) <EOL> MAX_NUM_CHARS_LOGGED_PER_PARAM = <NUM_LIT:1000> <EOL> def default_thrift_client_replace_if ( exception ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( exception , TException ) : <EOL> if ( isinstance ( exception , TApplicationException ) or <EOL> hasattr ( exception , '<STR_LIT>' ) ) : <EOL> return False <EOL> return True <EOL> class DummyStatsdClient : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> pass <EOL> def increment ( self , stats , sample_rate = <NUM_LIT:1> , tags = { } ) : <EOL> pass <EOL> def timing ( self , * args , ** kwargs ) : <EOL> pass <EOL> dummy_statsd = DummyStatsdClient ( ) <EOL> class PooledThriftClientMixinMetaclass ( type ) : <EOL> """<STR_LIT>""" <EOL> def __new__ ( mcs , classname , bases , classdict ) : <EOL> _thrift_client_cls = bases [ <NUM_LIT:0> ] <EOL> def _ensure_use_pooled_client ( self , method , timeout = <NUM_LIT> , <EOL> expiration = <NUM_LIT> , e2e_timeout = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> @ functools . wraps ( method ) <EOL> def wrap_with_pool ( * args , ** kwargs ) : <EOL> if self . client_pool . qsize ( ) == <NUM_LIT:0> : <EOL> log . info ( "<STR_LIT>" % <EOL> self . client_pool . pool_name ) <EOL> with self . client_pool . get_connection ( <EOL> timeout = timeout / MILLIS_PER_SEC , <EOL> expiration = expiration / MILLIS_PER_SEC , <EOL> replace_if = self . conn_replace_policy , <EOL> e2e_timeout = e2e_timeout / MILLIS_PER_SEC ) as client : <EOL> return method ( client , * args , ** kwargs ) <EOL> return wrap_with_pool <EOL> def __init__ ( self , host_provider , pool_size = <NUM_LIT:5> , timeout = <NUM_LIT> , <EOL> retry_policy = None , connection_wait_timeout = <NUM_LIT> , <EOL> connection_expiration = <NUM_LIT> , statsd_client = dummy_statsd , <EOL> socket_connection_timeout = None , <EOL> always_retry_on_new_host = False , <EOL> retry_count = <NUM_LIT:3> , <EOL> conn_replace_policy = default_thrift_client_replace_if , <EOL> protocol_factory = DEFAULT_THRIFT_PROTOCOL_FACTORY , <EOL> e2e_timeout = <NUM_LIT> , <EOL> is_ssl = False , validate = True , ca_certs = None , <EOL> failed_retry_policy_log_sample_rate = <NUM_LIT:1.0> ) : <EOL> """<STR_LIT>""" <EOL> class _ThriftClientMixinClass ( _thrift_client_cls , <EOL> ThriftClientMixin ) : <EOL> def get_connection_exception_class ( mixin_self ) : <EOL> return self . get_connection_exception_class ( ) <EOL> def _close_conn ( client ) : <EOL> """<STR_LIT>""" <EOL> client . teardown_connection ( ) <EOL> for method_name , method in getmembers ( _ThriftClientMixinClass , <EOL> predicate = ismethod ) : <EOL> if method_name [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] == '<STR_LIT:_>' : <EOL> continue <EOL> client_attr = getattr ( _thrift_client_cls , method_name , None ) <EOL> if client_attr is None or not ismethod ( client_attr ) : <EOL> continue <EOL> setattr ( self , method_name , <EOL> self . _ensure_use_pooled_client ( <EOL> method , <EOL> timeout = connection_wait_timeout , <EOL> expiration = connection_expiration , <EOL> e2e_timeout = e2e_timeout ) ) <EOL> pool_name = "<STR_LIT>" % ( _thrift_client_cls . __module__ , <EOL> _thrift_client_cls . __name__ ) <EOL> self . client_pool = ConnectionPool ( <EOL> pool_name = pool_name , <EOL> pool_size = pool_size , <EOL> close_conn_f = _close_conn , <EOL> conn_cls = _ThriftClientMixinClass , <EOL> host_provider = host_provider , <EOL> timeout = timeout , <EOL> statsd_client = statsd_client , <EOL> retry_policy = retry_policy , <EOL> socket_connection_timeout = socket_connection_timeout , <EOL> always_retry_on_new_host = always_retry_on_new_host , <EOL> retry_count = retry_count , <EOL> protocol_factory = protocol_factory , <EOL> is_ssl = is_ssl , <EOL> validate = validate , <EOL> ca_certs = ca_certs , <EOL> failed_retry_policy_log_sample_rate = failed_retry_policy_log_sample_rate ) <EOL> self . conn_replace_policy = conn_replace_policy <EOL> classdict [ '<STR_LIT>' ] = _ensure_use_pooled_client <EOL> classdict [ '<STR_LIT>' ] = __init__ <EOL> return super ( PooledThriftClientMixinMetaclass , mcs ) . __new__ ( <EOL> mcs , classname , bases , classdict ) <EOL> class PooledThriftClientMixin ( object ) : <EOL> """<STR_LIT>""" <EOL> __metaclass__ = PooledThriftClientMixinMetaclass <EOL> def _param_to_str ( param ) : <EOL> """<STR_LIT>""" <EOL> param_str = "<STR_LIT:%s>" % param <EOL> if len ( param_str ) > MAX_NUM_CHARS_LOGGED_PER_PARAM : <EOL> param_str = param_str [ : MAX_NUM_CHARS_LOGGED_PER_PARAM ] + "<STR_LIT>" <EOL> return param_str <EOL> def args_to_str ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> handler_args = "<STR_LIT:U+002C>" . join ( [ _param_to_str ( x ) for x in args ] ) <EOL> if kwargs : <EOL> kw_args = "<STR_LIT:U+002C>" . join ( <EOL> [ "<STR_LIT>" % ( k , _param_to_str ( v ) ) for ( k , v ) in kwargs . iteritems ( ) ] ) <EOL> handler_args = "<STR_LIT>" % ( handler_args , kw_args ) <EOL> if isinstance ( handler_args , unicode ) : <EOL> return handler_args . encode ( '<STR_LIT:utf-8>' , '<STR_LIT:ignore>' ) <EOL> return handler_args <EOL> except : <EOL> return '<STR_LIT>' <EOL> def _is_application_exception ( e ) : <EOL> """<STR_LIT>""" <EOL> return ( not isinstance ( e , TTransportException ) and <EOL> not isinstance ( e , ThriftConnectionTimeoutError ) and <EOL> not isinstance ( e , IOError ) ) <EOL> def _is_rpc_timeout ( e ) : <EOL> """<STR_LIT>""" <EOL> return isinstance ( e , socket . timeout ) <EOL> def ensure_connection ( service_name , method_name , method ) : <EOL> """<STR_LIT>""" <EOL> @ functools . wraps ( method ) <EOL> def method_wrapper ( self , * args , ** kwargs ) : <EOL> req_timeout_ms = kwargs . pop ( '<STR_LIT>' , self . timeout ) <EOL> conn_timeout_ms = kwargs . pop ( '<STR_LIT>' , <EOL> self . socket_connection_timeout ) <EOL> if conn_timeout_ms is None : <EOL> conn_timeout_ms = req_timeout_ms <EOL> retries_left = self . retry_count <EOL> while retries_left : <EOL> start_time = datetime . datetime . now ( ) <EOL> try : <EOL> try : <EOL> self . connect ( conn_timeout_ms , req_timeout_ms ) <EOL> except socket . timeout : <EOL> raise ThriftConnectionTimeoutError ( ) <EOL> result = method ( self . _client , * args , ** kwargs ) <EOL> time_taken = datetime . datetime . now ( ) - start_time <EOL> time_taken_ms = time_taken . total_seconds ( ) * <NUM_LIT:1000> <EOL> self . statsd_client . timing ( <EOL> "<STR_LIT>" . format ( service_name , method_name ) , <EOL> time_taken_ms , sample_rate = <NUM_LIT> ) <EOL> self . refresh_connection_if_needed ( ) <EOL> return result <EOL> except TApplicationException as e : <EOL> handler_args = args_to_str ( * args , ** kwargs ) <EOL> time_taken = datetime . datetime . now ( ) - start_time <EOL> time_taken_ms = time_taken . total_seconds ( ) * <NUM_LIT:1000> <EOL> log . info ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % ( <EOL> method_name , handler_args , self . host , <EOL> self . port , time_taken_ms , e ) ) <EOL> raise <EOL> except Exception as e : <EOL> t , v , tb = sys . exc_info ( ) <EOL> retries_left -= <NUM_LIT:1> <EOL> handler_args = args_to_str ( * args , ** kwargs ) <EOL> time_taken = datetime . datetime . now ( ) - start_time <EOL> time_taken_ms = time_taken . total_seconds ( ) * <NUM_LIT:1000> <EOL> if _is_application_exception ( e ) : <EOL> retry_policy_to_apply = self . retry_policy <EOL> if not retry_policy_to_apply : <EOL> retry_policy_to_apply = DEFAULT_RETRY_POLICY <EOL> if not retry_policy_to_apply . should_retry ( e ) : <EOL> if random . random ( ) < self . failed_retry_policy_log_sample_rate : <EOL> log . info ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % ( <EOL> method_name , handler_args , self . host , <EOL> self . port , time_taken_ms , e ) ) <EOL> raise t , v , tb <EOL> elif _is_rpc_timeout ( e ) : <EOL> log . info ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % ( <EOL> method_name , handler_args , self . host , <EOL> self . port , time_taken_ms , e ) ) <EOL> self . statsd_client . increment ( <EOL> "<STR_LIT>" , <EOL> sample_rate = <NUM_LIT> , <EOL> tags = { '<STR_LIT>' : self . host } ) <EOL> self . teardown_connection ( select_new_host = True ) <EOL> else : <EOL> log . info ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" , <EOL> method_name , handler_args , self . host , <EOL> self . port , retries_left , time_taken_ms , e ) <EOL> self . statsd_client . increment ( <EOL> "<STR_LIT>" , <EOL> sample_rate = <NUM_LIT> , <EOL> tags = { '<STR_LIT>' : self . host } ) <EOL> if self . always_retry_on_new_host or retries_left == <NUM_LIT:1> : <EOL> self . teardown_connection ( select_new_host = True ) <EOL> else : <EOL> self . teardown_connection ( select_new_host = False ) <EOL> if not retries_left : <EOL> log . info ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % ( <EOL> method_name , handler_args , self . host , self . port , <EOL> time_taken_ms , e ) ) <EOL> self . statsd_client . increment ( <EOL> "<STR_LIT>" , <EOL> sample_rate = <NUM_LIT> , <EOL> tags = { '<STR_LIT>' : self . host } ) <EOL> raise self . get_connection_exception_class ( ) ( e ) , None , tb <EOL> return method_wrapper <EOL> class ThriftClientConnectionEnsured ( type ) : <EOL> """<STR_LIT>""" <EOL> def __new__ ( cls , classname , bases , classdict ) : <EOL> _client_cls = bases [ <NUM_LIT:0> ] <EOL> def __init__ ( self , host_provider , timeout = <NUM_LIT> , <EOL> retry_policy = None , statsd_client = dummy_statsd , <EOL> socket_connection_timeout = None , <EOL> always_retry_on_new_host = False , <EOL> retry_count = <NUM_LIT:3> , <EOL> protocol_factory = DEFAULT_THRIFT_PROTOCOL_FACTORY , <EOL> is_ssl = False , validate = True , ca_certs = None , <EOL> failed_retry_policy_log_sample_rate = <NUM_LIT:1.0> ) : <EOL> """<STR_LIT>""" <EOL> self . _host_selector = RandomHostSelector ( <EOL> host_provider , expire_time = <NUM_LIT:0> , retry_time = <NUM_LIT:0> ) <EOL> host_port_pair = self . _host_selector . get_host ( ) <EOL> host , port = host_port_pair . split ( "<STR_LIT::>" ) <EOL> self . host = host <EOL> self . port = int ( port ) <EOL> self . timeout = timeout <EOL> self . is_ssl = is_ssl <EOL> self . validate = validate <EOL> self . ca_certs = ca_certs <EOL> self . retry_policy = retry_policy <EOL> self . retry_count = retry_count <EOL> self . socket_connection_timeout = socket_connection_timeout <EOL> self . always_retry_on_new_host = always_retry_on_new_host <EOL> self . connected = False <EOL> self . connected_at = <NUM_LIT:0> <EOL> self . requests_served = <NUM_LIT:0> <EOL> self . statsd_client = statsd_client <EOL> self . protocol_factory = protocol_factory <EOL> self . failed_retry_policy_log_sample_rate = failed_retry_policy_log_sample_rate <EOL> classdict [ '<STR_LIT>' ] = __init__ <EOL> service_name = "<STR_LIT>" <EOL> if _client_cls . __module__ : <EOL> service_name = _client_cls . __module__ <EOL> if service_name . split ( '<STR_LIT:.>' ) > <NUM_LIT:0> : <EOL> service_name = service_name . split ( '<STR_LIT:.>' ) [ - <NUM_LIT:1> ] <EOL> for method_name , method in getmembers ( _client_cls , predicate = ismethod ) : <EOL> if method_name [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] == '<STR_LIT:_>' : <EOL> continue <EOL> classdict [ method_name ] = ensure_connection ( <EOL> service_name , method_name , method ) <EOL> return type . __new__ ( cls , classname , bases , classdict ) <EOL> class ThriftClientMixin ( object ) : <EOL> """<STR_LIT>""" <EOL> __metaclass__ = ThriftClientConnectionEnsured <EOL> def connect ( self , conn_timeout_ms , req_timeout_ms ) : <EOL> """<STR_LIT>""" <EOL> if self . connected : <EOL> self . _socket . setTimeout ( req_timeout_ms ) <EOL> return <EOL> if self . is_ssl : <EOL> self . _socket = TSSLSocket ( self . host , self . port , self . validate , self . ca_certs ) <EOL> else : <EOL> self . _socket = TNoDelaySocket ( self . host , self . port ) <EOL> self . _socket . setTimeout ( conn_timeout_ms ) <EOL> self . _transport = TBufferedTransport ( self . _socket ) <EOL> self . _transport . open ( ) <EOL> self . _protocol = self . protocol_factory . getProtocol ( self . _transport ) <EOL> self . _client = self . __class__ . __bases__ [ <NUM_LIT:0> ] ( self . _protocol ) <EOL> self . connected = True <EOL> self . connected_at = time . time ( ) <EOL> self . _socket . setTimeout ( req_timeout_ms ) <EOL> def teardown_connection ( self , select_new_host = False ) : <EOL> """<STR_LIT>""" <EOL> if self . connected : <EOL> self . _transport . close ( ) <EOL> self . connected = False <EOL> self . connected_at = <NUM_LIT:0> <EOL> self . _transport = None <EOL> self . _socket = None <EOL> self . requests_served = <NUM_LIT:0> <EOL> if select_new_host : <EOL> new_host , new_port = self . _host_selector . get_host ( ) . split ( "<STR_LIT::>" ) <EOL> self . host = new_host <EOL> self . port = int ( new_port ) <EOL> def _is_connection_old ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . connected : <EOL> return False <EOL> return time . time ( ) - self . connected_at > SECS_FOR_CONNECTION_TEARDOWN <EOL> def refresh_connection_if_needed ( <EOL> self , num_requests = NUM_REQUESTS_FOR_CONNECTION_TEARDOWN ) : <EOL> """<STR_LIT>""" <EOL> self . requests_served += <NUM_LIT:1> <EOL> if ( self . _is_connection_old ( ) or self . requests_served == <EOL> NUM_REQUESTS_FOR_CONNECTION_TEARDOWN ) : <EOL> self . teardown_connection ( select_new_host = True ) </s>
