class Controller(object): def __init__(self, str_scenario_name, **kwargs): self.__initialize(str_scenario_name, **kwargs) def __enter__(self, str_scenario_name, **kwargs): self.__initialize(str_scenario_name, **kwargs) def __exit__(self, exc_type, exc_val, exc_tb): click.secho("Controller closed", fg="red") def __del__(self): click.secho("Controller closed", fg="red") def __initialize(self, str_scenario_name, **kwargs): # Self address self.ip_address = kwargs.get("ip_address", socket.gethostname()) # Controller port self.port = kwargs.get("port", 7878) # Allocate an empty dict to collect spawners data self.spawners = {} list_spawners = kwargs.get("spawners", [self.ip_address]) for str_spawner in list_spawners: self.spawners[str_spawner] = {} # Create a sender for each spawner self.spawners[str_spawner]["sender"] = Sender(str_spawner, self.port, "TCP") # Spawning mode (threads or subprocesses), defaults to threads self.spawn_mode = kwargs.get("spawn_mode", "thread") # Local flag to run everything on localhost (to be used in the future...), defaults to False self.local = kwargs.get("local", False) # Placeholder for scenario self.scenario = None # Scenario class name self.scenario_name = str_scenario_name # Scenario folder self.scenario_folder = utils.get_abs_path(kwargs.get("scenario_folder", os.getcwd())) # Scenario file, defaults to 'scenario.py' self.scenario_file = kwargs.get("scenario_file_path", "./scenario.py") # Scenario file path, defaults to standard scenario file self.scenario_file_path = utils.get_abs_path(self.scenario_file, self.scenario_folder) # Scenario duration, o be filled later self.scenario_duration = None # Result file, defaults to 'results/results.sqlite' file self.results_file = kwargs.get("results_file", "./results/results.sqlite") # Result file path self.results_file_path = utils.get_abs_path(self.results_file, self.scenario_folder) # Sysmonitor thread self.sysmonitor = Sysmonitor(self.ip_address, self.port, str_host_type="controller", bool_debug=False) # Log Collector thread self.logcollector = LogCollectorThread(self.results_file_path, self.ip_address, self.port) # Elapsed time since scenario duration self.scenario_elapsed_time = 0 def __load_scenario(self): # Get scenario from path and name self.scenario = utils.import_from_path( self.scenario_file_path, self.scenario_name, {"scenario_folder": self.scenario_folder} ) # Load navigations self.scenario.navigations_definition() # Fill scenario duration self.scenario_duration = self.scenario.get_scenario_duration() def __scale_ramps(self): # Get spawners number int_spawners_num = len(self.spawners) int_max_spawns = self.scenario.get_max_scenario_spawns() # Cycle through spawners scenarios and rescale max spawn number to match total spawn number for str_spawner_ip in self.spawners.iterkeys(): int_spawn_quota = int(round(int_max_spawns / int_spawners_num, 0)) dbl_spawn_quota = int_spawn_quota / int_max_spawns int_max_spawns -= int_spawn_quota int_spawners_num -= 1 self.spawners[str_spawner_ip]["spawn_quota"] = dbl_spawn_quota def __zip__scenario_folder(self): # Create a in-memory string file and write Zip file in it obj_in_memory_zip = StringIO() obj_zipfile = zipfile.ZipFile(obj_in_memory_zip, "w", zipfile.ZIP_DEFLATED) # Walk through files and folders for root, dirs, files in os.walk(self.scenario_folder): for filename in files: str_relpath = os.path.relpath(root, self.scenario_folder) obj_zipfile.write(os.path.join(root, filename), os.path.join(str_relpath, filename)) obj_zipfile.close() bin_archive = obj_in_memory_zip.getvalue() self.scenario_folder_encoded_zip = base64.b64encode(bin_archive) def __send_scenario(self): # Cycle through spawners and send serialized scenario class for str_spawner_ip in self.spawners.iterkeys(): dic_payload = { "scenarioBase64ZippedFolder": self.scenario_folder_encoded_zip, "spawnQuota": self.spawners[str_spawner_ip]["spawn_quota"], "controllerPort": self.port, "scenarioName": self.scenario_name, "scenarioFile": self.scenario_file, "resultsFile": self.results_file, } self.spawners[str_spawner_ip]["sender"].send("setup", dic_payload) def setup_scenario(self): self.__load_scenario() self.__scale_ramps() self.__zip__scenario_folder() self.__send_scenario() def start_scenario(self): # Start Logcollector and Sysmonitor threads self.logcollector.start() self.sysmonitor.start() # Cycle through spawners and send serialized scenario class dic_payload = {"spawnMode": self.spawn_mode} for str_spawner_ip in self.spawners.iterkeys(): self.spawners[str_spawner_ip]["sender"].send("start", dic_payload) # Wait until completion while self.scenario_elapsed_time <= self.scenario_duration: time.sleep(1) self.scenario_elapsed_time += 1 # Terminate Sysmonitor and Log Collector at the ennd of the scenario self.sysmonitor.terminate() self.sysmonitor.join() self.logcollector.terminate() self.logcollector.join() def shutdown_remotes(self): for str_spawner_ip in self.spawners.iterkeys(): self.spawners[str_spawner_ip]["sender"].send("shutdown", {})
class Spawner(StoppableThread): def __init__(self, str_server_address, int_controller_port, str_scenario_folder, str_scenario_name, str_scenario_file_path, str_results_file_path, dbl_spawn_quota): # Initialize the father stoppable thread class super(Spawner, self).__init__() # Instantiates everything self.__initialize(controller_ip_address=str_server_address, controller_port=int_controller_port, scenario_folder=str_scenario_folder, scenario_file=str_scenario_file_path, scenario_name=str_scenario_name, results_file=str_results_file_path, spawn_quota=dbl_spawn_quota) self.__set_scenario_class() self.__arm() def run(self): self.__run() def __initialize(self, **kwargs): # Controller and Log collector socket port self.port = kwargs.get('controller_port', 7878) # Controller address, void at the beginning self.server_address = kwargs.get('controller_ip_address', None) # Scenario folder, defaults to current directory self.scenario_folder = kwargs.get('scenario_folder', os.getcwd()) # Scenario file name, defaults to './scenario.py' self.scenario_file_path = utils.get_abs_path(kwargs.get('scenario_file', './scenario.py'), self.scenario_folder) # Results file path, defaults to './results/results.sqlite' self.results_file_path = utils.get_abs_path(kwargs.get('results_file', './results/results.sqlite'), self.scenario_folder) # Scenario name, defaults to 'Scenario' self.scenario_name = kwargs.get('scenario_name', 'Scenario') # Spawn quota for local spawner self.spawn_quota = kwargs.get('spawn_quota', 1) # Placeholder for scenario self.scenario = None # Armed flag, used to start and stop running self.armed = False # Elapsed time since scenario beginning self.elapsed_time = 0 # Internal message sender placeholder self.sender = Sender(self.server_address, self.port, 'UDP') if self.server_address else None # Sender spawn message poll interval self.sender_spawn_polling_interval = kwargs.get('sender_spawn_polling_interval', 5) self.sender_spawn_elapsed_time = 0.0 # Internal Sysmonitor object and related thread self.sysmonitor_polling_interval = kwargs.get('sysmonitor_polling_interval', 1.0) click.secho(utils.logify('Setting up Sysmonitor... ', 'Spawner')) self.sysmonitor = Sysmonitor(self.server_address, self.port, self.sysmonitor_polling_interval) click.secho(utils.logify('Sysmonitor started'), fg='green', bold=True) def __set_scenario_class(self): # Load scenario from temporary folder self.scenario = utils.import_from_path(self.scenario_file_path, self.scenario_name, {'scenario_folder': self.scenario_folder}) self.scenario.rescale_spawns_by_factor(self.spawn_quota) def __arm(self): # Arm itself self.armed = True def __unarm(self): # Disarm itself self.armed = False def __run(self): click.secho(utils.logify('Starting Sysmonitor... ', 'Spawner'), nl=False) self.sysmonitor.start() click.secho('DONE', fg='green', bold=True) self.scenario.configure() self.scenario.navigations_definition() self.scenario.scenario_start = utils.get_timestamp(False) self.scenario.scenario_duration =\ self.scenario.get_scenario_duration() list_navigations = self.scenario.get_navigation_names() # Update sender spawn message elapsed time self.sender_spawn_elapsed_time += self.elapsed_time # Write beginning message click.secho(utils.logify('Running scenario...', 'Spawner'), fg='green') while self.elapsed_time <= self.scenario.scenario_duration and self.armed: # Get elapsed time time_elapsed = utils.get_timestamp(False) - self.scenario.scenario_start self.elapsed_time = time_elapsed.total_seconds() # Cycle through navigations for str_navigation_name in list_navigations: # Get planned and current spawn number self.scenario.navigations[str_navigation_name]['planned_spawns'] =\ self.scenario.get_planned_spawns(str_navigation_name, self.elapsed_time) self.scenario.navigations[str_navigation_name]['current_spawns'] =\ len(self.scenario.navigations[str_navigation_name]['threads']) # Send spawn message, but only if is passed enough time from last sending if self.sender_spawn_elapsed_time >= self.sender_spawn_polling_interval: self.sender.send('spawns', { 'hostName': utils.get_ip_address(), 'timestamp': utils.get_timestamp(), 'navigationName': str_navigation_name, 'plannedSpawns': self.scenario.navigations[str_navigation_name]['planned_spawns'], 'runningSpawns': self.scenario.navigations[str_navigation_name]['current_spawns'] }) self.sender_spawn_elapsed_time = 0.0 int_spawns_difference = self.scenario.navigations[str_navigation_name][ 'current_spawns'] - self.scenario.navigations[ str_navigation_name]['planned_spawns'] # If there are less spawns than planned, add some spawns if int_spawns_difference < 0: for int_counter in range(0, -int_spawns_difference): str_navigation_path = self.scenario.get_navigation_path(str_navigation_name) str_id = utils.random_id(16) print(str_id) obj_spawn = Spawn(str_id, str_navigation_name, str_navigation_path, self.scenario_folder, self.server_address, self.port, self.scenario.settings) obj_spawn.start() self.scenario.navigations[str_navigation_name]['threads'].append(obj_spawn) # If there are more spawns than planned, start killing older spawns elif int_spawns_difference > 0: for int_counter in range(0, int_spawns_difference): obj_spawn = self.scenario.navigations[str_navigation_name]['threads'].pop(0) obj_spawn.terminate() obj_spawn.join() print('Killed one') self.__unarm() # Clean all the remaining threads click.secho(utils.logify('Cleaning up... ', 'Spawner'), nl=False) for str_navigation_name in list_navigations: for int_counter in range(0, len(self.scenario.navigations[str_navigation_name]['threads'])): obj_spawn = self.scenario.navigations[str_navigation_name]['threads'].pop(0) obj_spawn.terminate() obj_spawn.join() print('Killed one - Cleanup') click.secho('DONE', fg='green', bold=True) # Close Sysmonitor thread click.secho(utils.logify('Closing Sysmonitor... ', 'Spawner'), nl=False) self.sysmonitor.terminate() self.sysmonitor.join() click.secho('DONE', fg='green', bold=True) click.secho(utils.logify('Scenario gracefully completed!', 'Spawner'), fg='green', bold=True)