<s> from __future__ import absolute_import <EOL> import time <EOL> import socket <EOL> import struct <EOL> import logging <EOL> from . version import __version__ <EOL> try : <EOL> import ssl <EOL> except ImportError : <EOL> ssl = None <EOL> try : <EOL> from . snappy_socket import SnappySocket <EOL> except ImportError : <EOL> SnappySocket = None <EOL> try : <EOL> import simplejson as json <EOL> except ImportError : <EOL> import json <EOL> import tornado . iostream <EOL> import tornado . ioloop <EOL> try : <EOL> from tornado . simple_httpclient import _default_ca_certs as default_ca_certs <EOL> except ImportError : <EOL> from tornado . simple_httpclient import _DEFAULT_CA_CERTS <EOL> def default_ca_certs ( ) : <EOL> return _DEFAULT_CA_CERTS <EOL> from nsq import event , protocol <EOL> from . deflate_socket import DeflateSocket <EOL> logger = logging . getLogger ( __name__ ) <EOL> INIT = '<STR_LIT>' <EOL> DISCONNECTED = '<STR_LIT>' <EOL> CONNECTING = '<STR_LIT>' <EOL> CONNECTED = '<STR_LIT>' <EOL> DEFAULT_USER_AGENT = '<STR_LIT>' % __version__ <EOL> class AsyncConn ( event . EventedMixin ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( <EOL> self , <EOL> host , <EOL> port , <EOL> timeout = <NUM_LIT:1.0> , <EOL> heartbeat_interval = <NUM_LIT:30> , <EOL> requeue_delay = <NUM_LIT> , <EOL> tls_v1 = False , <EOL> tls_options = None , <EOL> snappy = False , <EOL> deflate = False , <EOL> deflate_level = <NUM_LIT:6> , <EOL> user_agent = DEFAULT_USER_AGENT , <EOL> output_buffer_size = <NUM_LIT:16> * <NUM_LIT> , <EOL> output_buffer_timeout = <NUM_LIT> , <EOL> sample_rate = <NUM_LIT:0> , <EOL> io_loop = None , <EOL> auth_secret = None , <EOL> msg_timeout = None ) : <EOL> assert isinstance ( host , ( str , unicode ) ) <EOL> assert isinstance ( port , int ) <EOL> assert isinstance ( timeout , float ) <EOL> assert isinstance ( tls_options , ( dict , None . __class__ ) ) <EOL> assert isinstance ( deflate_level , int ) <EOL> assert isinstance ( heartbeat_interval , int ) and heartbeat_interval >= <NUM_LIT:1> <EOL> assert isinstance ( requeue_delay , int ) and requeue_delay >= <NUM_LIT:0> <EOL> assert isinstance ( output_buffer_size , int ) and output_buffer_size >= <NUM_LIT:0> <EOL> assert isinstance ( output_buffer_timeout , int ) and output_buffer_timeout >= <NUM_LIT:0> <EOL> assert isinstance ( sample_rate , int ) and sample_rate >= <NUM_LIT:0> and sample_rate < <NUM_LIT:100> <EOL> assert isinstance ( auth_secret , ( str , unicode , None . __class__ ) ) <EOL> assert tls_v1 and ssl or not tls_v1 , '<STR_LIT>' <EOL> assert msg_timeout is None or ( isinstance ( msg_timeout , ( float , int ) ) and msg_timeout > <NUM_LIT:0> ) <EOL> self . state = INIT <EOL> self . host = host <EOL> self . port = port <EOL> self . timeout = timeout <EOL> self . last_recv_timestamp = time . time ( ) <EOL> self . last_msg_timestamp = time . time ( ) <EOL> self . in_flight = <NUM_LIT:0> <EOL> self . rdy = <NUM_LIT:0> <EOL> self . rdy_timeout = None <EOL> self . max_rdy_count = <NUM_LIT> <EOL> self . tls_v1 = tls_v1 <EOL> self . tls_options = tls_options <EOL> self . snappy = snappy <EOL> self . deflate = deflate <EOL> self . deflate_level = deflate_level <EOL> self . hostname = socket . gethostname ( ) <EOL> self . short_hostname = self . hostname . split ( '<STR_LIT:.>' ) [ <NUM_LIT:0> ] <EOL> self . heartbeat_interval = heartbeat_interval * <NUM_LIT:1000> <EOL> self . msg_timeout = int ( msg_timeout * <NUM_LIT:1000> ) if msg_timeout else None <EOL> self . requeue_delay = requeue_delay <EOL> self . io_loop = io_loop <EOL> if not self . io_loop : <EOL> self . io_loop = tornado . ioloop . IOLoop . instance ( ) <EOL> self . output_buffer_size = output_buffer_size <EOL> self . output_buffer_timeout = output_buffer_timeout <EOL> self . sample_rate = sample_rate <EOL> self . user_agent = user_agent <EOL> self . _authentication_required = False <EOL> self . auth_secret = auth_secret <EOL> self . socket = None <EOL> self . stream = None <EOL> self . _features_to_enable = [ ] <EOL> self . last_rdy = <NUM_LIT:0> <EOL> self . rdy = <NUM_LIT:0> <EOL> self . callback_queue = [ ] <EOL> super ( AsyncConn , self ) . __init__ ( ) <EOL> @ property <EOL> def id ( self ) : <EOL> return str ( self ) <EOL> def __str__ ( self ) : <EOL> return self . host + '<STR_LIT::>' + str ( self . port ) <EOL> def connect ( self ) : <EOL> if self . state not in [ INIT , DISCONNECTED ] : <EOL> return <EOL> self . socket = socket . socket ( socket . AF_INET , socket . SOCK_STREAM ) <EOL> self . socket . settimeout ( self . timeout ) <EOL> self . socket . setblocking ( <NUM_LIT:0> ) <EOL> self . stream = tornado . iostream . IOStream ( self . socket , io_loop = self . io_loop ) <EOL> self . stream . set_close_callback ( self . _socket_close ) <EOL> self . state = CONNECTING <EOL> self . on ( event . CONNECT , self . _on_connect ) <EOL> self . on ( event . DATA , self . _on_data ) <EOL> self . stream . connect ( ( self . host , self . port ) , self . _connect_callback ) <EOL> def _connect_callback ( self ) : <EOL> self . state = CONNECTED <EOL> self . stream . write ( protocol . MAGIC_V2 ) <EOL> self . _start_read ( ) <EOL> self . trigger ( event . CONNECT , conn = self ) <EOL> def _read_bytes ( self , size , callback ) : <EOL> try : <EOL> self . stream . read_bytes ( size , callback ) <EOL> except IOError : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . ConnectionClosedError ( '<STR_LIT>' ) , <EOL> ) <EOL> def _start_read ( self ) : <EOL> self . _read_bytes ( <NUM_LIT:4> , self . _read_size ) <EOL> def _socket_close ( self ) : <EOL> self . state = DISCONNECTED <EOL> self . trigger ( event . CLOSE , conn = self ) <EOL> def close ( self ) : <EOL> self . stream . close ( ) <EOL> def _read_size ( self , data ) : <EOL> try : <EOL> size = struct . unpack ( '<STR_LIT>' , data ) [ <NUM_LIT:0> ] <EOL> except Exception : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . IntegrityError ( '<STR_LIT>' ) , <EOL> ) <EOL> return <EOL> self . _read_bytes ( size , self . _read_body ) <EOL> def _read_body ( self , data ) : <EOL> try : <EOL> self . trigger ( event . DATA , conn = self , data = data ) <EOL> except Exception : <EOL> logger . exception ( '<STR_LIT>' ) <EOL> self . io_loop . add_callback ( self . _start_read ) <EOL> def send ( self , data ) : <EOL> self . stream . write ( data ) <EOL> def upgrade_to_tls ( self , options = None ) : <EOL> assert ssl , '<STR_LIT>' <EOL> self . io_loop . remove_handler ( self . socket . fileno ( ) ) <EOL> opts = { <EOL> '<STR_LIT>' : ssl . CERT_REQUIRED , <EOL> '<STR_LIT>' : default_ca_certs ( ) <EOL> } <EOL> opts . update ( options or { } ) <EOL> self . socket = ssl . wrap_socket ( self . socket , ssl_version = ssl . PROTOCOL_TLSv1 , <EOL> do_handshake_on_connect = False , ** opts ) <EOL> self . stream = tornado . iostream . SSLIOStream ( self . socket , io_loop = self . io_loop ) <EOL> self . stream . set_close_callback ( self . _socket_close ) <EOL> self . stream . _do_ssl_handshake ( ) <EOL> def upgrade_to_snappy ( self ) : <EOL> assert SnappySocket , '<STR_LIT>' <EOL> existing_data = self . stream . _consume ( self . stream . _read_buffer_size ) <EOL> self . socket = SnappySocket ( self . socket ) <EOL> self . socket . bootstrap ( existing_data ) <EOL> self . stream . socket = self . socket <EOL> def upgrade_to_deflate ( self ) : <EOL> existing_data = self . stream . _consume ( self . stream . _read_buffer_size ) <EOL> self . socket = DeflateSocket ( self . socket , self . deflate_level ) <EOL> self . socket . bootstrap ( existing_data ) <EOL> self . stream . socket = self . socket <EOL> def send_rdy ( self , value ) : <EOL> try : <EOL> self . send ( protocol . ready ( value ) ) <EOL> except Exception , e : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . SendError ( '<STR_LIT>' % value , e ) , <EOL> ) <EOL> return False <EOL> self . last_rdy = value <EOL> self . rdy = value <EOL> return True <EOL> def _on_connect ( self , ** kwargs ) : <EOL> identify_data = { <EOL> '<STR_LIT>' : self . short_hostname , <EOL> '<STR_LIT>' : self . hostname , <EOL> '<STR_LIT>' : self . heartbeat_interval , <EOL> '<STR_LIT>' : True , <EOL> '<STR_LIT>' : self . tls_v1 , <EOL> '<STR_LIT>' : self . snappy , <EOL> '<STR_LIT>' : self . deflate , <EOL> '<STR_LIT>' : self . deflate_level , <EOL> '<STR_LIT>' : self . output_buffer_timeout , <EOL> '<STR_LIT>' : self . output_buffer_size , <EOL> '<STR_LIT>' : self . sample_rate , <EOL> '<STR_LIT>' : self . user_agent <EOL> } <EOL> if self . msg_timeout : <EOL> identify_data [ '<STR_LIT>' ] = self . msg_timeout <EOL> self . trigger ( event . IDENTIFY , conn = self , data = identify_data ) <EOL> self . on ( event . RESPONSE , self . _on_identify_response ) <EOL> try : <EOL> self . send ( protocol . identify ( identify_data ) ) <EOL> except Exception , e : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . SendError ( '<STR_LIT>' , e ) , <EOL> ) <EOL> def _on_identify_response ( self , data , ** kwargs ) : <EOL> self . off ( event . RESPONSE , self . _on_identify_response ) <EOL> if data == '<STR_LIT:OK>' : <EOL> logger . warning ( '<STR_LIT>' ) <EOL> return self . trigger ( event . READY , conn = self ) <EOL> try : <EOL> data = json . loads ( data ) <EOL> except ValueError : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . IntegrityError ( <EOL> '<STR_LIT>' % <EOL> data <EOL> ) , <EOL> ) <EOL> return <EOL> self . trigger ( event . IDENTIFY_RESPONSE , conn = self , data = data ) <EOL> if self . tls_v1 and data . get ( '<STR_LIT>' ) : <EOL> self . _features_to_enable . append ( '<STR_LIT>' ) <EOL> if self . snappy and data . get ( '<STR_LIT>' ) : <EOL> self . _features_to_enable . append ( '<STR_LIT>' ) <EOL> if self . deflate and data . get ( '<STR_LIT>' ) : <EOL> self . _features_to_enable . append ( '<STR_LIT>' ) <EOL> if data . get ( '<STR_LIT>' ) : <EOL> self . _authentication_required = True <EOL> if data . get ( '<STR_LIT>' ) : <EOL> self . max_rdy_count = data . get ( '<STR_LIT>' ) <EOL> else : <EOL> logger . warn ( '<STR_LIT>' ) <EOL> self . max_rdy_count = <NUM_LIT> <EOL> self . on ( event . RESPONSE , self . _on_response_continue ) <EOL> self . _on_response_continue ( conn = self , data = None ) <EOL> def _on_response_continue ( self , data , ** kwargs ) : <EOL> if self . _features_to_enable : <EOL> feature = self . _features_to_enable . pop ( <NUM_LIT:0> ) <EOL> if feature == '<STR_LIT>' : <EOL> self . upgrade_to_tls ( self . tls_options ) <EOL> elif feature == '<STR_LIT>' : <EOL> self . upgrade_to_snappy ( ) <EOL> elif feature == '<STR_LIT>' : <EOL> self . upgrade_to_deflate ( ) <EOL> return <EOL> self . off ( event . RESPONSE , self . _on_response_continue ) <EOL> if self . auth_secret and self . _authentication_required : <EOL> self . on ( event . RESPONSE , self . _on_auth_response ) <EOL> self . trigger ( event . AUTH , conn = self , data = self . auth_secret ) <EOL> try : <EOL> self . send ( protocol . auth ( self . auth_secret ) ) <EOL> except Exception , e : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . SendError ( '<STR_LIT>' , e ) , <EOL> ) <EOL> return <EOL> self . trigger ( event . READY , conn = self ) <EOL> def _on_auth_response ( self , data , ** kwargs ) : <EOL> try : <EOL> data = json . loads ( data ) <EOL> except ValueError : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . IntegrityError ( <EOL> '<STR_LIT>' % data <EOL> ) , <EOL> ) <EOL> return <EOL> self . off ( event . RESPONSE , self . _on_auth_response ) <EOL> self . trigger ( event . AUTH_RESPONSE , conn = self , data = data ) <EOL> return self . trigger ( event . READY , conn = self ) <EOL> def _on_data ( self , data , ** kwargs ) : <EOL> self . last_recv_timestamp = time . time ( ) <EOL> frame , data = protocol . unpack_response ( data ) <EOL> if frame == protocol . FRAME_TYPE_MESSAGE : <EOL> self . last_msg_timestamp = time . time ( ) <EOL> self . rdy = max ( self . rdy - <NUM_LIT:1> , <NUM_LIT:0> ) <EOL> self . in_flight += <NUM_LIT:1> <EOL> message = protocol . decode_message ( data ) <EOL> message . on ( event . FINISH , self . _on_message_finish ) <EOL> message . on ( event . REQUEUE , self . _on_message_requeue ) <EOL> message . on ( event . TOUCH , self . _on_message_touch ) <EOL> self . trigger ( event . MESSAGE , conn = self , message = message ) <EOL> elif frame == protocol . FRAME_TYPE_RESPONSE and data == '<STR_LIT>' : <EOL> self . send ( protocol . nop ( ) ) <EOL> self . trigger ( event . HEARTBEAT , conn = self ) <EOL> elif frame == protocol . FRAME_TYPE_RESPONSE : <EOL> self . trigger ( event . RESPONSE , conn = self , data = data ) <EOL> elif frame == protocol . FRAME_TYPE_ERROR : <EOL> self . trigger ( event . ERROR , conn = self , error = protocol . Error ( data ) ) <EOL> def _on_message_requeue ( self , message , backoff = True , time_ms = - <NUM_LIT:1> , ** kwargs ) : <EOL> if backoff : <EOL> self . trigger ( event . BACKOFF , conn = self ) <EOL> else : <EOL> self . trigger ( event . CONTINUE , conn = self ) <EOL> self . in_flight -= <NUM_LIT:1> <EOL> try : <EOL> time_ms = self . requeue_delay * message . attempts * <NUM_LIT:1000> if time_ms < <NUM_LIT:0> else time_ms <EOL> self . send ( protocol . requeue ( message . id , time_ms ) ) <EOL> except Exception , e : <EOL> self . close ( ) <EOL> self . trigger ( event . ERROR , conn = self , error = protocol . SendError ( <EOL> '<STR_LIT>' % ( message . id , time_ms ) , e ) ) <EOL> def _on_message_finish ( self , message , ** kwargs ) : <EOL> self . trigger ( event . RESUME , conn = self ) <EOL> self . in_flight -= <NUM_LIT:1> <EOL> try : <EOL> self . send ( protocol . finish ( message . id ) ) <EOL> except Exception , e : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . SendError ( '<STR_LIT>' % message . id , e ) , <EOL> ) <EOL> def _on_message_touch ( self , message , ** kwargs ) : <EOL> try : <EOL> self . send ( protocol . touch ( message . id ) ) <EOL> except Exception , e : <EOL> self . close ( ) <EOL> self . trigger ( <EOL> event . ERROR , <EOL> conn = self , <EOL> error = protocol . SendError ( '<STR_LIT>' % message . id , e ) , <EOL> ) </s>
