<s> import pynitro <EOL> import diesel <EOL> from functools import partial <EOL> from diesel import log , loglevels <EOL> from diesel . events import Waiter , StopWaitDispatch <EOL> from diesel . util . queue import Queue <EOL> from diesel . util . event import Event <EOL> class DieselNitroSocket ( Waiter ) : <EOL> def __init__ ( self , bind = None , connect = None , ** kwargs ) : <EOL> Waiter . __init__ ( self ) <EOL> self . destroyed = False <EOL> kwargs [ '<STR_LIT>' ] = <NUM_LIT:1> <EOL> self . socket = pynitro . NitroSocket ( ** kwargs ) <EOL> self . _early_value = None <EOL> from diesel . runtime import current_app <EOL> from diesel . hub import IntWrap <EOL> if bind : <EOL> assert not connect <EOL> self . socket . bind ( bind ) <EOL> elif connect : <EOL> assert not bind <EOL> self . socket . connect ( connect ) <EOL> self . hub = current_app . hub <EOL> self . fd = IntWrap ( self . socket . fileno ( ) ) <EOL> self . read_gate = Event ( ) <EOL> self . hub . register ( self . fd , self . messages_exist , self . error , self . error ) <EOL> self . sent = <NUM_LIT:0> <EOL> self . received = <NUM_LIT:0> <EOL> def _send_op ( self , op ) : <EOL> while True : <EOL> try : <EOL> op ( ) <EOL> except pynitro . NitroFull : <EOL> diesel . sleep ( <NUM_LIT> ) <EOL> else : <EOL> self . sent += <NUM_LIT:1> <EOL> return <EOL> def recv ( self ) : <EOL> while True : <EOL> try : <EOL> m = self . socket . recv ( self . socket . NOWAIT ) <EOL> except pynitro . NitroEmpty : <EOL> self . read_gate . clear ( ) <EOL> self . read_gate . wait ( ) <EOL> else : <EOL> self . received += <NUM_LIT:1> <EOL> return m <EOL> def send ( self , frame , flags = <NUM_LIT:0> ) : <EOL> return self . _send_op ( <EOL> partial ( self . socket . send , frame , self . socket . NOWAIT | flags ) ) <EOL> def reply ( self , orig , frame , flags = <NUM_LIT:0> ) : <EOL> return self . _send_op ( <EOL> partial ( self . socket . reply , orig , frame , self . socket . NOWAIT | flags ) ) <EOL> def process_fire ( self , dc ) : <EOL> if not self . _early_value : <EOL> got = self . ready_early ( ) <EOL> if not got : <EOL> raise StopWaitDispatch ( ) <EOL> assert self . _early_value <EOL> v = self . _early_value <EOL> self . _early_value = None <EOL> return v <EOL> def ready_early ( self ) : <EOL> if self . _early_value : <EOL> return True <EOL> try : <EOL> m = self . socket . recv ( self . socket . NOWAIT ) <EOL> except pynitro . NitroEmpty : <EOL> self . read_gate . clear ( ) <EOL> return False <EOL> else : <EOL> self . received += <NUM_LIT:1> <EOL> self . _early_value = m <EOL> return True <EOL> def messages_exist ( self ) : <EOL> '''<STR_LIT>''' <EOL> self . read_gate . set ( ) <EOL> diesel . fire ( self ) <EOL> def error ( self ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> def destroy ( self ) : <EOL> if not self . destroyed : <EOL> self . hub . unregister ( self . fd ) <EOL> del self . socket <EOL> self . destroyed = True <EOL> def __enter__ ( self ) : <EOL> return self <EOL> def __del__ ( self ) : <EOL> self . destroy ( ) <EOL> def __exit__ ( self , * args ) : <EOL> self . destroy ( ) <EOL> class DieselNitroService ( object ) : <EOL> """<STR_LIT>""" <EOL> name = '<STR_LIT>' <EOL> default_log_level = loglevels . DEBUG <EOL> timeout = <NUM_LIT:10> <EOL> def __init__ ( self , uri , logger = None , log_level = None ) : <EOL> self . uri = uri <EOL> self . nitro_socket = None <EOL> self . log = logger or None <EOL> self . selected_log_level = log_level <EOL> self . clients = { } <EOL> self . outgoing = Queue ( ) <EOL> self . incoming = Queue ( ) <EOL> self . name = self . name or self . __class__ . __name__ <EOL> self . _incoming_loop = None <EOL> try : <EOL> self . should_run = True <EOL> except AttributeError : <EOL> pass <EOL> if self . log and self . selected_log_level is not None : <EOL> self . selected_log_level = None <EOL> warnings . warn ( <EOL> "<STR_LIT>" , <EOL> RuntimeWarning , <EOL> stacklevel = <NUM_LIT:2> , <EOL> ) <EOL> def _create_server_socket ( self ) : <EOL> self . nitro_socket = DieselNitroSocket ( bind = self . uri ) <EOL> def _setup_the_logging_system ( self ) : <EOL> if not self . log : <EOL> if self . selected_log_level is not None : <EOL> log_level = self . selected_log_level <EOL> else : <EOL> log_level = self . default_log_level <EOL> log_name = self . name or self . __class__ . __name__ <EOL> self . log = log . name ( log_name ) <EOL> self . log . min_level = log_level <EOL> def _handle_client_requests_and_responses ( self , remote_client ) : <EOL> assert self . nitro_socket <EOL> queues = [ remote_client . incoming ] <EOL> try : <EOL> while True : <EOL> ( evt , value ) = diesel . first ( waits = queues , sleep = self . timeout ) <EOL> if evt is remote_client . incoming : <EOL> assert isinstance ( value , Message ) <EOL> remote_client . async_frame = value . orig_frame <EOL> resp = self . handle_client_packet ( value . data , remote_client . context ) <EOL> if resp : <EOL> if isinstance ( resp , basestring ) : <EOL> output = [ resp ] <EOL> else : <EOL> output = iter ( resp ) <EOL> for part in output : <EOL> msg = Message ( <EOL> value . orig_frame , <EOL> remote_client . identity , <EOL> self . serialize_message ( remote_client . identity , part ) , <EOL> ) <EOL> self . outgoing . put ( msg ) <EOL> elif evt == '<STR_LIT>' : <EOL> break <EOL> finally : <EOL> self . _cleanup_client ( remote_client ) <EOL> def _cleanup_client ( self , remote_client ) : <EOL> del self . clients [ remote_client . identity ] <EOL> self . cleanup_client ( remote_client ) <EOL> self . log . debug ( "<STR_LIT>" % remote_client . identity ) <EOL> def _handle_all_inbound_and_outbound_traffic ( self ) : <EOL> assert self . nitro_socket <EOL> queues = [ self . nitro_socket , self . outgoing ] <EOL> socket = self . nitro_socket <EOL> make_frame = pynitro . NitroFrame <EOL> while self . should_run : <EOL> ( queue , msg ) = diesel . first ( waits = queues ) <EOL> if queue is self . outgoing : <EOL> socket . reply ( msg . orig_frame , make_frame ( msg . data ) ) <EOL> else : <EOL> id , obj = self . parse_message ( msg . data ) <EOL> msg . clear_data ( ) <EOL> msg = Message ( msg , id , obj ) <EOL> if msg . identity not in self . clients : <EOL> self . _register_client ( msg ) <EOL> self . clients [ msg . identity ] . incoming . put ( msg ) <EOL> def _register_client ( self , msg ) : <EOL> remote = RemoteClient . from_message ( msg ) <EOL> self . clients [ msg . identity ] = remote <EOL> self . register_client ( remote , msg ) <EOL> diesel . fork_child ( self . _handle_client_requests_and_responses , remote ) <EOL> def __call__ ( self ) : <EOL> return self . run ( ) <EOL> def run ( self ) : <EOL> self . _create_server_socket ( ) <EOL> self . _setup_the_logging_system ( ) <EOL> self . _handle_all_inbound_and_outbound_traffic ( ) <EOL> def handle_client_packet ( self , packet , context ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def cleanup_client ( self , remote_client ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> def register_client ( self , remote_client , msg ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> def parse_message ( self , raw_data ) : <EOL> """<STR_LIT>""" <EOL> return None , raw_data <EOL> def serialize_message ( self , identity , raw_data ) : <EOL> """<STR_LIT>""" <EOL> return raw_data <EOL> def async_send ( self , identity , msg ) : <EOL> """<STR_LIT>""" <EOL> remote_client = self . clients [ identity ] <EOL> out = self . serialize_message ( msg ) <EOL> self . outgoing . put ( <EOL> Message ( <EOL> remote_client . async_frame , <EOL> identity , <EOL> out ) ) <EOL> class RemoteClient ( object ) : <EOL> def __init__ ( self , identity ) : <EOL> self . identity = identity <EOL> self . incoming = Queue ( ) <EOL> self . context = { } <EOL> self . async_frame = None <EOL> @ classmethod <EOL> def from_message ( cls , msg ) : <EOL> return cls ( msg . identity ) <EOL> class Message ( object ) : <EOL> def __init__ ( self , frame , identity , data ) : <EOL> self . orig_frame = frame <EOL> self . identity = identity <EOL> self . data = data </s>
