Exemple #1
0
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", {})
Exemple #2
0
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)