def wait_for_request(self): # type: () -> Tuple[bytes, List[bytes]] """Wait for a REQUEST command from the broker and return the client address and message body frames. Will internally handle timeouts, heartbeats and check for protocol errors and disconnect commands. """ command, frames = self._receive() if command == p.DISCONNECT: self._log.debug("Got DISCONNECT from broker; Disconnecting") self._disconnect() raise error.Disconnected("Disconnected on message from broker") elif command != p.REQUEST: raise error.ProtocolError( "Unexpected message type from broker: {}".format( command.decode('utf8'))) if len(frames) < 3: raise error.ProtocolError( "Unexpected REQUEST message size, got {} frames, expecting at least 3" .format(len(frames))) client_addr = frames[0] request = frames[2:] return client_addr, request
def _parse_incoming_message(frames): # type: (List[bytes]) -> protocol.Message """Parse and verify incoming message """ if len(frames) < 4: raise error.ProtocolError( "Unexpected message length: expecting at least 4 frames, got {}" .format(len(frames))) if frames[1] != b'': raise error.ProtocolError( "Expecting empty frame 1, got {} bytes".format(len(frames[1]))) return protocol.Message(client=frames[0], header=frames[2], command=frames[3], message=frames[4:])
def handle_message(self, message): # type: (protocol.Message) -> None """Handle incoming message """ if message.header == protocol.WORKER_HEADER: self._handle_worker_message(message) elif message.header == protocol.CLIENT_HEADER: self._handle_client_message(message) else: raise error.ProtocolError("Unexpected protocol header: {}".format( message.header)) self._dispatch_queued_messages()
def _verify_message(message): # type: (List[bytes]) -> Tuple[bytes, List[bytes]] if len(message) < 3: raise error.ProtocolError( "Unexpected message length, expecting at least 3 frames, got {}" .format(len(message))) if message.pop(0) != b'': raise error.ProtocolError( "Expecting first message frame to be empty") if message[0] != p.WORKER_HEADER: print(message) raise error.ProtocolError( "Unexpected protocol header [{}], expecting [{}]".format( message[0].decode('utf8'), p.WORKER_HEADER.decode('utf8'))) if message[1] not in {p.DISCONNECT, p.HEARTBEAT, p.REQUEST}: raise error.ProtocolError( "Unexpected message type [{}], expecting either HEARTBEAT, REQUEST or " "DISCONNECT".format(message[1].decode('utf8'))) return message[1], message[2:]
def _handle_client_message(self, message): # type: (protocol.Message) -> None """Handle message from a client """ assert message.command == protocol.REQUEST if len(message.message) < 2: raise error.ProtocolError( "Client REQUEST message is expected to be at least 2 frames long, got {}" .format(len(message.message))) service_name = message.message[0] body = message.message[1:] # TODO: Plug-in MMA handling self._log.debug("Queueing client request from %d to %s", id_to_int(message.client), service_name.decode('ascii')) self._services.queue_client_request(message.client, service_name, body)
def _handle_worker_message(self, message): # type: (protocol.Message) -> None """Handle message from a worker """ if message.command == protocol.HEARTBEAT: self._handle_worker_heartbeat(message.client) elif message.command == protocol.READY: self._handle_worker_ready(message.client, message.message[0]) elif message.command == protocol.FINAL: self._handle_worker_final(message.client, message.message[0], message.message[2:]) elif message.command == protocol.PARTIAL: self._handle_worker_partial(message.client, message.message[0], message.message[2:]) elif message.command == protocol.DISCONNECT: self._handle_worker_disconnect(message.client) else: self._send_to_worker(message.client, protocol.DISCONNECT) raise error.ProtocolError( "Unexpected command from worker: {}".format(message.command))