boss.boss - heros-boss 0.5.1 documentation

Skip to content
Back to top

# Source code for boss.boss

    from .helper import log
    from .configuration import WorkerConfigurationDocument
    import asyncio
    from collections.abc import Callable
    from concurrent.futures import ProcessPoolExecutor

    class BOSS:
        hero_config_sources: list

        def __init__(
            self,
            name: str,
            configs: list,
            loop: asyncio.AbstractEventLoop,
            pool: ProcessPoolExecutor,
            realm: str,
            expose: bool = True,
        ):
            """
            A BOSS object.

            Args:
                name: name of the Boss object
                loop: asyncio event loop that can be used to schedule tasks
                configs: List of WorkerConfigurationDocuments to specify the workers
                realm: Name of the realm the BOSS starts their HEROs in.
                expose: Specify whether the BOSS object itself should be exposed as HERO
            """
            self.name = name
            self.heros: dict = {}
            self.hero_config_sources: list = []
            self._loop = loop
            self._pool = pool
            self.realm = realm

            self.start_all()

        def _config_from_name(self, name):
            if name not in self.heros:
                raise NameError(f"HERO with name {name} not run by this BOSS. Available HEROs are {self.heros.keys()}")
            return self.heros[name]["config"]

        def add_hero_source(self, parser: Callable[[str], WorkerConfigurationDocument], target: str):
            """
            Adds a data source to the BOSS from which HERO configurations are loaded.

            Args:
                parser: A function which takes :code:`target` as an argument and returns a
                    :py:class:`boss.configuration.WorkerConfigurationDocument` dict with the HERO config information.
                target: Target for the :code:`parser`, for example an URL for
                    :py:meth:`boss.configuration.WorkerConfigurationDocument.parse_url`
            """
            self.hero_config_sources.append([parser, target])

        def refresh_hero_sources(self, auto_start: bool = True):
            """
            Refresh the HERO configuration data from the registered sources, updating existing HEROs, adding new ones,
            and removing those that are no longer present in the sources.
            HEROs that were running before are only restarted if their configuration changed. HEROs without configuration
            changes stay untouched.

            Args:
                auto_start: If True, automatically starts a *new* hero source if its configuration is
                           loaded. Defaults to True. Does not influence the behavior of HEROs that are already registered.

            """

            current_heros = list(self.heros.keys())

            def handle_hero_dict(config: dict):
                if config is not None:
                    if "rows" in config:
                        for conf_row in config["rows"]:
                            if "doc" in conf_row:
                                handle_hero_dict(conf_row["doc"])
                            else:
                                handle_hero_dict(conf_row)
                    elif "_id" in config:
                        name = config["_id"]
                        do_start = False
                        if name in self.heros:
                            new_conf_doc = WorkerConfigurationDocument(config)
                            new_conf_doc.generate_tags(self)
                            if self.heros[name]["config"] != new_conf_doc:
                                # config changed, we need to restart
                                if self.heros[name]["status"] == "running":
                                    do_start = True  # was running before, auto start with the reloaded config
                                    self.stop_hero(name)
                                self.heros[name]["config"] = new_conf_doc
                            current_heros.remove(name)
                        else:
                            do_start = auto_start
                        self.add_hero(config, auto_start=do_start)

            for parser, target in self.hero_config_sources:
                log.info("refreshing HERO source %s", target)
                config = parser(target)
                handle_hero_dict(config)

            for deleted_hero in current_heros:
                # hero not found in reloaded sources, delete
                log.info("Removing HERO %s from BOSS", deleted_hero)
                self.remove_hero(deleted_hero)

        def add_hero(self, config: WorkerConfigurationDocument | dict, auto_start: bool = True):
            """
            Start a new HERO and keep it running. Note that the id of the HERO specified in the config must be unique.

            Args:
                config: configuration for the new HERO. If a dict is given, it is converted into
                    a WorkerConfigurationDocument.
                auto_start: If true the new HERO is immediately started after adding
            """

            if not isinstance(config, dict):
                return None
            if not isinstance(config, WorkerConfigurationDocument):
                config = WorkerConfigurationDocument(config)

            name = config["_id"]

            if name not in self.heros:
                self.heros[name] = {"config": config, "object": None, "status": "stopped"}

            if auto_start:
                self.start_hero(name)

        def start_hero(self, name):
            """
            Start HERO with given name.
            """
            config = self._config_from_name(name)
            status = self.heros[name]["status"]

            if status not in "running":
                try:
                    self.heros[name]["object"] = config.build_hero_for_boss(self, self.realm)
                    log.info(f"creating HERO with name {config['_id']} from class {config['classname']}")

                    # calling setup hook if it exists
                    if hasattr(self.heros[config["_id"]], "_setup") and callable(
                        getattr(self.heros[name]["object"], "_setup")
                    ):
                        self.heros[name]["object"]._setup()

                    self.heros[name]["status"] = "running"

                except KeyError as e:
                    log.debug("", exc_info=True)
                    log.error(f"creating HERO with invalid dict: {config} failed:  {e}")
                except Exception as e:
                    log.debug("", exc_info=True)
                    log.error(f"creating HERO with name {config['_id']} from class {config['classname']} failed: {e}")

        def stop_hero(self, name: str):
            config = self._config_from_name(name)
            status = self.heros[name]["status"]

            if status == "running":
                log.info(f"destroying HERO with name {config['_id']} from class {config['classname']}")
                obj = self.heros[name]["object"]

                # calling teardown hook if it exists
                if hasattr(obj, "_teardown") and callable(getattr(obj, "_teardown")):
                    obj._teardown()

                obj._destroy_hero()
                del self.heros[name]["object"]
                self.heros[name]["object"] = None
                del obj

                self.heros[name]["status"] = "stopped"

        def restart_hero(self, name: str):
            self.stop_hero(name)
            self.start_hero(name)

        def remove_hero(self, name: str):
            self.stop_hero(name)

            if name in self.heros:
                del self.heros[name]

        def start_all(self):
            for hero_name in self.heros.keys():
                self.start_hero(hero_name)

        def stop_all(self):
            for hero_name in self.heros.keys():
                self.stop_hero(hero_name)

        def status(self):
            return self.heros