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