def _end_session(self): """ Terminates the injection session for all connected hosts """ msg_end = MessageBuilder.command_session(time(), end=True) self._client.broadcast_msg(msg_end) session_closed = 0 session_sent = self._client.get_n_registered_hosts() session_check_start = time() session_check_now = time() while session_check_now - session_check_start < self._sessionWait and session_closed < session_sent: # We wait until we have received an ack for the termination from all of the connected hosts, or we time out if self._client.peek_msg_queue() > 0: addr, msg = self._client.pop_msg_queue() if msg[MessageBuilder.FIELD_TYPE] == MessageBuilder.ACK_YES: InjectorController.logger.info( "Injection session closed with engine %s" % formatipport(addr)) if not self._suppressOutput: self._writers[addr].write_entry( MessageBuilder.command_session( msg[MessageBuilder.FIELD_TIME], end=True)) session_closed += 1 else: # If we receive a message that is not an ack after all tasks have terminated, something is wrong InjectorController.logger.error( "Ack expected from engine %s, got %s" % (formatipport(addr), msg[MessageBuilder.FIELD_TYPE])) sleep(self._sleepPeriod) session_check_now = time() # All of the execution log writers are closed, and the session finishes if not self._suppressOutput: for writer in self._writers.values(): writer.close()
def _process_msg_inject(self, addr, msg): """ Processes incoming message for clients involved in an injection session :param addr: The address of the sender :param msg: The message dictionary """ # We process status messages for connections that are in the queue is_status, status = MessageClient.is_status_message(msg) if is_status and status == MessageClient.CONNECTION_LOST_MSG: # If connection has been lost with an host, we remove its pendingTasks entry if not self._suppressOutput: self._writers[addr].write_entry( MessageBuilder.status_connection(time())) elif is_status and status == MessageClient.CONNECTION_RESTORED_MSG: # If connection has been restored with an host, we send a new session start command self._client.send_msg( addr, MessageBuilder.command_session(self._session_id)) self._client.send_msg( addr, MessageBuilder.command_set_time(self._get_timestamp(time()))) elif is_status and status == MessageClient.CONNECTION_FINALIZED_MSG: self._pendingTasks.pop(addr, None) # If all connections to servers were finalized we assume that the injection can be terminated if len(self._pendingTasks) == 0: self._endReached = True self._reader.close() else: msg_type = msg[MessageBuilder.FIELD_TYPE] if msg_type != MessageBuilder.ACK_YES and msg_type != MessageBuilder.ACK_NO: # Ack messages are not written to the output log if not self._suppressOutput: self._writers[addr].write_entry(msg) # We log on the terminal the content of the message in a pretty form if msg_type == MessageBuilder.STATUS_START: InjectorController.logger.info( "Task %s started on host %s" % (msg[MessageBuilder.FIELD_DATA], formatipport(addr))) elif msg_type == MessageBuilder.STATUS_RESTART: InjectorController.logger.info( "Task %s restarted on host %s" % (msg[MessageBuilder.FIELD_DATA], formatipport(addr))) elif msg_type == MessageBuilder.STATUS_END: InjectorController.logger.info( "Task %s terminated successfully on host %s" % (msg[MessageBuilder.FIELD_DATA], formatipport(addr))) # If a task terminates, we remove its sequence number from the set of pending tasks for the host self._pendingTasks[addr].discard( msg[MessageBuilder.FIELD_SEQNUM]) if not self._suppressOutput: self._write_task_output(addr, msg) elif msg_type == MessageBuilder.STATUS_ERR: InjectorController.logger.error( "Task %s terminated with error code %s on host %s" % (msg[MessageBuilder.FIELD_DATA], str(msg[MessageBuilder.FIELD_ERR]), formatipport(addr))) self._pendingTasks[addr].discard( msg[MessageBuilder.FIELD_SEQNUM]) if not self._suppressOutput: self._write_task_output(addr, msg) elif msg_type == MessageBuilder.ACK_YES: # ACK messages after the initialization phase are received ONLY when a connection is restored, # and the session must be resumed InjectorController.logger.warning( "Session resumed with engine %s" % formatipport(addr)) # If the ack msg contains an error, it means all previously running tasks have been lost if not self._suppressOutput: self._writers[addr].write_entry( MessageBuilder.status_connection(time(), restored=True)) if MessageBuilder.FIELD_ERR in msg: self._pendingTasks[addr] = set() if not self._suppressOutput: self._writers[addr].write_entry( MessageBuilder.status_reset( msg[MessageBuilder.FIELD_TIME])) elif msg_type == MessageBuilder.ACK_NO: InjectorController.logger.warning( "Session cannot be resumed with engine %s" % formatipport(addr)) self._client.remove_host(addr)
def _init_session(self, workload_name): """ Initializes the injection session for all connected hosts :param workload_name: The name of the workload to be injected :return: the number of hosts that have accepted the injection start command, and the timestamp ID of the session """ session_start_timestamp = time() msg_start = MessageBuilder.command_session(session_start_timestamp) self._client.broadcast_msg(msg_start) self._writers = {} self._outputsDirs = {} self._pendingTasks = {} session_accepted = set() session_replied = 0 session_sent = self._client.get_n_registered_hosts() session_check_start = time() session_check_now = time() while session_check_now - session_check_start < self._sessionWait and session_replied < session_sent: # We wait until we receive an ack (positive or negative) from all connected hosts, or either we time out if self._client.peek_msg_queue() > 0: addr, msg = self._client.pop_msg_queue() if msg[MessageBuilder.FIELD_TYPE] == MessageBuilder.ACK_YES: # If an host replies to the injection start command with a positive ack, its log writer is # instantiated, together with its entry in the pendingTasks dictionary InjectorController.logger.info( "Injection session started with engine %s" % formatipport(addr)) session_accepted.add(addr) session_replied += 1 self._outputsDirs[addr] = format_output_directory( self._resultsDir, addr, workload_name) # The outputs directory needs to be flushed before starting the new injection session if not self._suppressOutput: if isdir(self._outputsDirs[addr]): rmtree(self._outputsDirs[addr], ignore_errors=True) self._writers[addr] = ExecutionLogWriter( format_injection_filename(self._resultsDir, addr, workload_name)) self._writers[addr].write_entry( MessageBuilder.command_session( msg[MessageBuilder.FIELD_TIME])) self._pendingTasks[addr] = set() elif msg[MessageBuilder.FIELD_TYPE] == MessageBuilder.ACK_NO: # If an host rejects the injection start command, we discard it InjectorController.logger.warning( "Injection session request rejected by engine %s" % formatipport(addr)) session_replied += 1 self._client.remove_host(addr) sleep(self._sleepPeriod) session_check_now = time() if session_check_now - session_check_start >= self._sessionWait: # If we have reached the time out, it means that not all of the connected hosts have replied. This is # highly unlikely, but could still happen. In this case, we remove all hosts that have not replied InjectorController.logger.warning( "Injection session startup reached the timeout limit") for addr in self._client.get_registered_hosts(): if addr not in session_accepted: self._client.remove_host(addr) return len(session_accepted), session_start_timestamp