class Agent(AgentBase): def __init__(self, token: str, dsp_host: str = 'localhost', dsp_port: int = 9999): logger.info('Starting Agent') super(Agent, self).__init__() self.socket = RequestConnection(dsp_host, dsp_port) self.broker = None self.token = token def __enter__(self): self.socket.establish() return self def __exit__(self, *exc_info): self.close() def close(self): if self.broker: self.broker.close() self.socket.close() def register(self): request = deepcopy(Register_agent) if self.name: request['name'] = self.name request['token'] = self.token reply = self.socket.send(request) if reply['result']: self.id = reply['id'] self.sync(reply) return reply['result'] def init_broker(self): """ Request Agent queues on Broker from Dispatcher. """ request = deepcopy(Agent_queues) request['token'] = self.token request['id'] = self.id reply = self.socket.send(request) if reply['result']: self.sync(reply) self.broker = Broker(reply['broker']['host']) self.broker.connect() self.broker.declare(reply['broker']['task'], reply['broker']['result']) def pulse(self) -> bool: request = deepcopy(Pulse) request['id'] = self.id reply = self.socket.send(request) self.sync(reply['reply']) return reply['result'] def sync(self, reply: dict): self.last_sync = datetime.utcnow() if 'commands' in reply: self.commands = reply['commands'] def apply_commands(self): for command in self.commands: _method = getattr(self, command) if not _method(): logger.error(f'Method "{_method}" has failed to apply') return False else: return True def disconnect(self): request = deepcopy(Disconnect) request['id'] = self.id reply = self.socket.send(request) if reply['result']: self.broker.close() self.broker = None self.id = 0 return reply['result'] def shutdown(self): self.__exit__() exit(0)
class Dispatcher: def __init__(self, ip: str = '*', port: Union[int, str] = '', broker_host: str = ''): logger.info('Starting Dispatcher') self.socket = ReplyConnection(ip, port) self.broker = Broker(broker_host if broker_host else ip) self.agents = {} self.request_handler = self.default_request_handler self._next_free_id = INIT_AGENT_ID self._listen = True self._interrupt: Union[None, Callable] = None def __enter__(self): self.socket.establish() self.configure_broker() return self def __exit__(self, *exc_info): logger.info(f'Closing Dispatcher connection:{self.socket}') self.socket.close() self.broker.close() def configure_broker(self): self.broker._inactivity_timeout = 0.01 * SECOND if self.broker.connect() and not self.broker.input_queue: self.broker.setup_exchange() self.broker.declare( input_queue=compose_queue(RoutingKeys.DISPATCHER)) def listen(self, polling_timeout: int = 10 * SECOND): while self._listen: expired = self.socket.listen(self.request_handler, polling_timeout) if self._interrupt and self._interrupt(expired): break if not self.broker.input_queue and not self.broker.channel.is_open: self.configure_broker() else: for task in self.broker.pulling_generator(): logger.debug(f'Got dispatcher task {task}') def default_request_handler(self, request: dict): commands = { Commands.Register_agent: self._register_agent_handler, Commands.Agent_queues: self._agent_queues_handler, Commands.Pulse: self._pulse_handler, Commands.Client_queues: self._client_handler, Commands.Relay: self._relay, Commands.Disconnect: self._disconnect_handler } assert request['command'] in commands, 'Command ' \ f'{request["command"]} is not registered in dispatcher request handler' command = commands[request['command']] return command(request) def _register_agent_handler(self, request: dict): logger.info( f'Registration request received {request["name"]}({self._next_free_id})' ) request['id'] = self._next_free_id agent = RemoteAgent(self._next_free_id) agent.name = request['name'] agent.token = request['token'] self.agents[self._next_free_id] = agent logger.info(f'New agent id={agent.id}') request['result'] = True self._next_free_id += 1 return request def _agent_queues_handler(self, request: dict): """ Returns Host and queues that agent should connect to. """ agent = self.agents[request['id']] logger.info(f'Agent queues request received from {agent}') config = Database.get_agent_param(agent.token) request['broker']['host'] = config['broker'] request['broker']['task'] = compose_queue(RoutingKeys.TASK) request['broker']['result'] = compose_queue(RoutingKeys.RESULTS) request['result'] = True return request def _pulse_handler(self, request: dict): logger.info(f'Pulse request received {request["id"]}') agent = self.agents[request['id']] reply = agent.sync(request) return reply def _client_handler(self, request: dict): logger.info(f'Client queues are requested by: {request["name"]}') config = Database.get_client_param(request['token']) request['broker']['host'] = config['broker'] request['broker']['task'] = compose_queue(RoutingKeys.TASK) request['broker']['result'] = compose_queue(request['name']) request['result'] = True return request def _disconnect_handler(self, request: dict): """ Removes agent instance on dispatcher. """ logger.info(f'Disconnect request received {request["id"]}') if request['id'] in self.agents: self.agents.pop(request['id']) request['result'] = True return request def _relay(self, request: dict): """ Returns received request. """ logger.debug(f'Relay request called') return request