<s> """<STR_LIT>""" <EOL> import os , sys <EOL> import collections <EOL> import errno <EOL> import json <EOL> import socket <EOL> import threading <EOL> import time <EOL> import zmq <EOL> from . import config <EOL> from . import core <EOL> from . import sockets <EOL> _logger = core . get_logger ( __name__ ) <EOL> Continue = object ( ) <EOL> Empty = object ( ) <EOL> def _unpack ( message ) : <EOL> return json . loads ( message . decode ( config . ENCODING ) ) <EOL> def _pack ( message ) : <EOL> return json . dumps ( message ) . encode ( config . ENCODING ) <EOL> def timed_out ( started_at , wait_for_s ) : <EOL> if wait_for_s is config . FOREVER : <EOL> return False <EOL> else : <EOL> return time . time ( ) > started_at + wait_for_s <EOL> def _bind_with_timeout ( bind_function , args , n_tries = <NUM_LIT:3> , retry_interval_s = <NUM_LIT:0.5> ) : <EOL> """<STR_LIT>""" <EOL> n_tries_left = n_tries <EOL> while n_tries_left > <NUM_LIT:0> : <EOL> try : <EOL> return bind_function ( * args ) <EOL> except zmq . error . ZMQError as exc : <EOL> _logger . warn ( "<STR_LIT>" , exc , n_tries_left ) <EOL> n_tries_left -= <NUM_LIT:1> <EOL> except OSError as exc : <EOL> if exc . errno == errno . EADDRINUSE : <EOL> _logger . warn ( "<STR_LIT>" , exc , n_tries_left ) <EOL> n_tries_left -= <NUM_LIT:1> <EOL> else : <EOL> raise <EOL> else : <EOL> raise core . SocketAlreadyExistsError ( "<STR_LIT>" % n_tries ) <EOL> class _Service ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , name , address ) : <EOL> self . name = name <EOL> self . address = address <EOL> self . advertise_at = <NUM_LIT:0> <EOL> def __str__ ( self ) : <EOL> return "<STR_LIT>" % ( self . name , self . address , time . ctime ( self . advertise_at ) ) <EOL> class _Command ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , action , params ) : <EOL> self . action = action <EOL> self . params = params <EOL> self . started_at = time . time ( ) <EOL> self . response = Empty <EOL> def __str__ ( self ) : <EOL> return "<STR_LIT>" % ( self . action , self . params , time . ctime ( self . started_at ) , self . response ) <EOL> class _Beacon ( threading . Thread ) : <EOL> """<STR_LIT>""" <EOL> rpc_port = <NUM_LIT> <EOL> beacon_port = <NUM_LIT> <EOL> finder_timeout_s = <NUM_LIT> <EOL> beacon_message_size = <NUM_LIT> <EOL> time_between_broadcasts_s = config . BEACON_ADVERT_FREQUENCY_S <EOL> def __init__ ( self ) : <EOL> threading . Thread . __init__ ( self ) <EOL> self . setDaemon ( True ) <EOL> self . _stop_event = threading . Event ( ) <EOL> self . _services_to_advertise = collections . deque ( ) <EOL> self . _services_found = { } <EOL> self . _command = None <EOL> self . socket = socket . socket ( socket . AF_INET , socket . SOCK_DGRAM , socket . IPPROTO_UDP ) <EOL> self . socket . setsockopt ( socket . SOL_SOCKET , socket . SO_BROADCAST , <NUM_LIT:1> ) <EOL> self . socket . bind ( ( "<STR_LIT>" , self . beacon_port ) ) <EOL> self . socket_fd = self . socket . fileno ( ) <EOL> self . poller = zmq . Poller ( ) <EOL> self . poller . register ( self . socket , zmq . POLLIN ) <EOL> self . rpc = sockets . context . socket ( zmq . REP ) <EOL> _bind_with_timeout ( self . rpc . bind , ( "<STR_LIT>" % self . rpc_port , ) ) <EOL> def stop ( self ) : <EOL> _logger . debug ( "<STR_LIT>" ) <EOL> self . _stop_event . set ( ) <EOL> def do_advertise ( self , started_at , name , address , fail_if_exists ) : <EOL> _logger . debug ( "<STR_LIT>" , name , address , fail_if_exists ) <EOL> canonical_address = core . address ( address ) <EOL> for service in self . _services_to_advertise : <EOL> if service . name == name : <EOL> if fail_if_exists : <EOL> _logger . error ( "<STR_LIT>" , name , service . address ) <EOL> return None <EOL> else : <EOL> _logger . warn ( "<STR_LIT>" , name , service . address ) <EOL> self . _services_to_advertise . append ( _Service ( name , canonical_address ) ) <EOL> self . _services_found [ name ] = canonical_address <EOL> return canonical_address <EOL> def do_discover ( self , started_at , name , wait_for_s ) : <EOL> _logger . debug ( "<STR_LIT>" , name , wait_for_s ) <EOL> discovered = self . _services_found . get ( name ) <EOL> if discovered : <EOL> return discovered <EOL> if timed_out ( started_at , wait_for_s ) : <EOL> return None <EOL> else : <EOL> return Continue <EOL> def do_discover_all ( self , started_at ) : <EOL> _logger . debug ( "<STR_LIT>" ) <EOL> return list ( self . _services_found . items ( ) ) <EOL> def do_reset ( self , started_at ) : <EOL> _logger . debug ( "<STR_LIT>" ) <EOL> self . _services_found . clear ( ) <EOL> self . _services_to_advertise . clear ( ) <EOL> def do_stop ( self , started_at ) : <EOL> _logger . debug ( "<STR_LIT>" ) <EOL> self . stop ( ) <EOL> def listen_for_one_advert ( self ) : <EOL> events = dict ( self . poller . poll ( <NUM_LIT:1000> * self . finder_timeout_s ) ) <EOL> if self . socket_fd not in events : <EOL> return <EOL> message , source = self . socket . recvfrom ( self . beacon_message_size ) <EOL> service_name , service_address = _unpack ( message ) <EOL> self . _services_found [ service_name ] = service_address <EOL> def broadcast_one_advert ( self ) : <EOL> if self . _services_to_advertise : <EOL> next_service = self . _services_to_advertise [ <NUM_LIT:0> ] <EOL> if next_service . advertise_at < time . time ( ) : <EOL> message = _pack ( [ next_service . name , next_service . address ] ) <EOL> self . socket . sendto ( message , <NUM_LIT:0> , ( "<STR_LIT>" , self . beacon_port ) ) <EOL> next_service . advertise_at = time . time ( ) + self . time_between_broadcasts_s <EOL> self . _services_to_advertise . rotate ( - <NUM_LIT:1> ) <EOL> def poll_command_request ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> message = self . rpc . recv ( zmq . NOBLOCK ) <EOL> except zmq . ZMQError as exc : <EOL> if exc . errno == zmq . EAGAIN : <EOL> return <EOL> else : <EOL> raise <EOL> _logger . debug ( "<STR_LIT>" , message ) <EOL> segments = _unpack ( message ) <EOL> action , params = segments [ <NUM_LIT:0> ] , segments [ <NUM_LIT:1> : ] <EOL> _logger . debug ( "<STR_LIT>" , action , params ) <EOL> self . _command = _Command ( action , params ) <EOL> def process_command ( self ) : <EOL> if not self . _command : <EOL> return <EOL> else : <EOL> _logger . debug ( "<STR_LIT>" , self . _command . action ) <EOL> command = self . _command <EOL> _logger . debug ( "<STR_LIT>" , self . _command . action , self . _command . params , self . _command . started_at ) <EOL> function = getattr ( self , "<STR_LIT>" + command . action . lower ( ) , None ) <EOL> if not function : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> else : <EOL> try : <EOL> result = function ( command . started_at , * command . params ) <EOL> except : <EOL> _logger . exception ( "<STR_LIT>" , command . action , command . params ) <EOL> result = None <EOL> _logger . debug ( "<STR_LIT>" , result ) <EOL> if result is Continue : <EOL> return <EOL> self . _command . response = result <EOL> def poll_command_reponse ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _command . response is not Empty : <EOL> _logger . debug ( "<STR_LIT>" , self . _command . response ) <EOL> self . rpc . send ( _pack ( self . _command . response ) ) <EOL> self . _command = None <EOL> def run ( self ) : <EOL> _logger . info ( "<STR_LIT>" ) <EOL> while not self . _stop_event . wait ( <NUM_LIT:0> ) : <EOL> try : <EOL> if not self . _command : <EOL> self . poll_command_request ( ) <EOL> self . broadcast_one_advert ( ) <EOL> self . listen_for_one_advert ( ) <EOL> if self . _command : <EOL> self . process_command ( ) <EOL> self . poll_command_reponse ( ) <EOL> except : <EOL> _logger . exception ( "<STR_LIT>" ) <EOL> break <EOL> _logger . info ( "<STR_LIT>" ) <EOL> self . rpc . close ( ) <EOL> self . socket . close ( ) <EOL> _beacon = None <EOL> _remote_beacon = object ( ) <EOL> def _start_beacon ( ) : <EOL> """<STR_LIT>""" <EOL> global _beacon <EOL> if _beacon is None : <EOL> _logger . debug ( "<STR_LIT>" ) <EOL> try : <EOL> _beacon = _Beacon ( ) <EOL> except OSError as exc : <EOL> if exc . errno == errno . EADDRINUSE : <EOL> _logger . warn ( "<STR_LIT>" ) <EOL> _beacon = _remote_beacon <EOL> else : <EOL> raise <EOL> else : <EOL> _beacon . start ( ) <EOL> def _rpc ( action , * args ) : <EOL> _logger . debug ( "<STR_LIT>" , action , args ) <EOL> with sockets . context . socket ( zmq . REQ ) as socket : <EOL> socket . connect ( "<STR_LIT>" % _Beacon . rpc_port ) <EOL> socket . send ( _pack ( [ action ] + list ( args ) ) ) <EOL> return _unpack ( socket . recv ( ) ) <EOL> def advertise ( name , address = None , fail_if_exists = False ) : <EOL> """<STR_LIT>""" <EOL> _start_beacon ( ) <EOL> return _rpc ( "<STR_LIT>" , name , address , fail_if_exists ) <EOL> return address <EOL> def discover ( name , wait_for_s = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> _start_beacon ( ) <EOL> t0 = time . time ( ) <EOL> while True : <EOL> discovery = _rpc ( "<STR_LIT>" , name , <NUM_LIT:0.5> ) <EOL> if discovery : <EOL> return discovery <EOL> if timed_out ( t0 , wait_for_s ) : <EOL> return None <EOL> def discover_all ( ) : <EOL> """<STR_LIT>""" <EOL> _start_beacon ( ) <EOL> return _rpc ( "<STR_LIT>" ) <EOL> def discover_group ( group , exclude = None ) : <EOL> """<STR_LIT>""" <EOL> _start_beacon ( ) <EOL> if exclude is None : <EOL> names_to_exclude = set ( ) <EOL> else : <EOL> names_to_exclude = set ( exclude ) <EOL> all_discovered = _rpc ( "<STR_LIT>" ) <EOL> return [ ( name , address ) <EOL> for ( name , address ) in all_discovered <EOL> if name . startswith ( "<STR_LIT>" % group ) <EOL> and name not in names_to_exclude <EOL> ] <EOL> def reset_beacon ( ) : <EOL> """<STR_LIT>""" <EOL> _start_beacon ( ) <EOL> return _rpc ( "<STR_LIT>" ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> _start_beacon ( ) </s>
