BB: Use threading to download logs

Reduces time to download logs by approx. 40%, but may be unnecessary
after resolving https://github.com/guardicore/monkey/issues/2383
This commit is contained in:
Mike Salvatore 2022-10-02 14:16:29 -04:00
parent e4155648c1
commit 6a783d9c3e
1 changed files with 20 additions and 6 deletions

View File

@ -1,6 +1,7 @@
import logging import logging
from pathlib import Path from pathlib import Path
from typing import Mapping, Sequence from threading import Thread
from typing import List, Mapping
from common.types import MachineID from common.types import MachineID
from monkey_island.cc.models import Agent, Machine from monkey_island.cc.models import Agent, Machine
@ -12,7 +13,7 @@ class MonkeyLogsDownloader(object):
def __init__(self, island_client, log_dir_path): def __init__(self, island_client, log_dir_path):
self.island_client = island_client self.island_client = island_client
self.log_dir_path = Path(log_dir_path) self.log_dir_path = Path(log_dir_path)
self.monkey_log_paths: Sequence[Path] = [] self.monkey_log_paths: List[Path] = []
def download_monkey_logs(self): def download_monkey_logs(self):
try: try:
@ -21,16 +22,29 @@ class MonkeyLogsDownloader(object):
agents = self.island_client.get_agents() agents = self.island_client.get_agents()
machines = self.island_client.get_machines() machines = self.island_client.get_machines()
download_threads: List[Thread] = []
# TODO: Does downloading logs concurrently still improve performance after resolving
# https://github.com/guardicore/monkey/issues/2383?
for agent in agents: for agent in agents:
log_file_path = self._get_log_file_path(agent, machines) t = Thread(target=self._download_log, args=(agent, machines), daemon=True)
log_contents = self.island_client.get_agent_log(agent.id) t.start()
download_threads.append(t)
MonkeyLogsDownloader._write_log_to_file(log_file_path, log_contents) for thread in download_threads:
thread.join()
self.monkey_log_paths.append(log_file_path)
except Exception as err: except Exception as err:
LOGGER.exception(err) LOGGER.exception(err)
def _download_log(self, agent: Agent, machines: Mapping[MachineID, Machine]):
log_file_path = self._get_log_file_path(agent, machines)
log_contents = self.island_client.get_agent_log(agent.id)
MonkeyLogsDownloader._write_log_to_file(log_file_path, log_contents)
self.monkey_log_paths.append(log_file_path)
def _get_log_file_path(self, agent: Agent, machines: Mapping[MachineID, Machine]) -> Path: def _get_log_file_path(self, agent: Agent, machines: Mapping[MachineID, Machine]) -> Path:
try: try:
machine_ip = machines[agent.machine_id].network_interfaces[0].ip machine_ip = machines[agent.machine_id].network_interfaces[0].ip