colony/colony/core/roles/webview/role.py

import threading
from concurrent.futures import ThreadPoolExecutor
from enum import Enum

from werkzeug.serving import make_server

from colony.core.roles.role import Role
from colony.core.roles.unit.role import Unit
from colony.core.roles.webview.app import create_app
from colony.settings import setting

VIEW_INTERVAL = 5
VIEW_TIMEOUT = 2
VIEW_MAX_AGE = 20
VIEW_WORKERS = 8


class WebView(Role):
    class COMMANDS(Enum):
        pass

    def __init__(self, context):
        host, port = setting("web.host"), setting("web.port")
        self.app = create_app(self._status, self._chatter)
        try:
            self._server = make_server(host, port, self.app, threaded=True)
        except SystemExit:
            raise OSError(f"Web view cannot listen on {host}:{port}") from None
        self.port = self._server.server_port

        self._views = {}
        self._views_lock = threading.Lock()

        super().__init__({}, context)

    def _unit(self):
        unit = self.context.self().get_role_instances().get("Unit")
        if unit is None:
            raise RuntimeError("The Unit role is not available")
        return unit

    def _chatter(self):
        return self.context.self().get_role_instances().get("Chatter")

    def _status(self, compact=False):
        payload = self._unit().show_cluster()
        own_key = next((member["ip"], member["port"]) for member in payload["members"] if member["self"])
        views = self._fresh_views()
        views[own_key] = payload["members"]
        if compact:
            known = {(member["ip"], member["port"]) for member in payload["members"]}
            views = {key: rows for key, rows in views.items() if key in known}
        payload["views"] = [list(key) for key in sorted(views)]
        if compact:
            payload["routes"] = self._routes(views, payload["members"])
        else:
            payload["links"] = self._links(views)
        return payload

    @staticmethod
    def _routes(views, members):
        """Per reporting node, only what differs from "reached directly": the relayed and the unreached members."""
        index = {(member["ip"], member["port"]): i for i, member in enumerate(members)}
        routes = []
        for source in sorted(views):
            if source not in index:
                continue
            rows = {(row["ip"], row["port"]): row for row in views[source] if "direct" in row and not row.get("self")}
            via, none = {}, []
            for target, i in index.items():
                if target == source or not members[i]["alive"]:
                    continue
                row = rows.get(target)
                if row is None or not row.get("alive") or not (row["direct"] or row.get("via")):
                    none.append(i)
                elif not row["direct"]:
                    gateway = index.get(tuple(row["via"]))
                    if gateway is None:
                        none.append(i)
                    else:
                        via.setdefault(str(gateway), []).append(i)
            if via or none:
                routes.append({"from": index[source], "via": via, "none": none})
        return routes

    @staticmethod
    def _links(views):
        links = []
        for source, rows in views.items():
            for row in rows:
                if row.get("self") or not row.get("alive") or "direct" not in row:
                    continue
                if row["direct"] or row.get("via"):
                    links.append({"from": list(source), "to": [row["ip"], row["port"]], "via": None if row["direct"] else row["via"]})
        return links

    def _fresh_views(self):
        now = self.context.membership().now()
        with self._views_lock:
            return {key: rows for key, (seen, rows) in self._views.items() if now - seen <= VIEW_MAX_AGE}

    def _fetch_view(self, key):
        try:
            resp = self.context.execute_to_unregistered_node(key[0], key[1], Unit.COMMANDS.SHOW_CLUSTER, timeout=VIEW_TIMEOUT)
        except Exception:
            return None

        result = resp.get("result") if isinstance(resp, dict) and resp.get("status") == "ok" else None
        rows = result.get("members") if isinstance(result, dict) else None
        if not isinstance(rows, list) or not all(isinstance(row, dict) for row in rows):
            return None
        return rows if any("direct" in row for row in rows) else None

    def refresh_views(self):
        targets = [
            (member["ip"], member["port"]) for member in self._unit().show_cluster()["members"]
            if member["alive"] and not member["self"]
        ]
        with ThreadPoolExecutor(max_workers=max(1, min(VIEW_WORKERS, len(targets)))) as pool:
            results = list(pool.map(self._fetch_view, targets))

        now = self.context.membership().now()
        with self._views_lock:
            for key, rows in zip(targets, results):
                if rows is not None:
                    self._views[key] = (now, rows)
            self._views = {key: entry for key, entry in self._views.items() if key in targets}

    def _refresh_loop(self):
        while not self._stop_event.is_set():
            try:
                self.refresh_views()
            except Exception:
                self.logger.exception("Could not refresh the topology")
            self._stop_event.wait(VIEW_INTERVAL)

    def start(self):
        with self._start_lock:
            if self._thread is not None:
                return
            self._thread = threading.Thread(target=self._server.serve_forever, daemon=True, name=str(self))
            self._thread.start()
            threading.Thread(target=self._refresh_loop, daemon=True, name="WebViewTopology").start()

        self.logger.info(f"Web view at http://{self.context.self().ip}:{self.port}")

    def stop(self):
        super().stop()
        with self._start_lock:
            started = self._thread is not None
        if started:
            self._server.shutdown()
        self._server.server_close()