ejabberd-tools/prometheus.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer
from socket import AF_INET6
from time import time
from urllib.parse import parse_qs, urlparse

from prometheus_client import CollectorRegistry, Gauge, generate_latest, CONTENT_TYPE_LATEST

from config import Config
from metrics import EjabberdMetrics


class DynamicMetricsHandler(BaseHTTPRequestHandler):
    """HTTP handler that gives metrics from ``core.REGISTRY``."""

    def do_GET(self):
        params = parse_qs(urlparse(self.path).query)
        registry = self.generator(params)
        if "name[]" in params:
            registry = registry.restricted_registry(params["name[]"])
        try:
            output = generate_latest(registry)
        except Exception as exception:
            self.send_error(500, f"error generating metric output: {exception}")
            raise
        self.send_response(200)
        self.send_header("Content-Type", CONTENT_TYPE_LATEST)
        self.end_headers()
        self.wfile.write(output)

    @staticmethod
    def factory(registry_generator):
        DynMetricsHandler = type("MetricsHandler", (DynamicMetricsHandler, object), {"generator": registry_generator})
        return DynMetricsHandler


class Prometheus:
    def __init__(self, metrics):
        self.ttl = 10
        self._last_update = 0
        self._metrics = metrics

    def handler(self, metrics_handler):
        now = time()
        if now >= (self._last_update + self.ttl):
            self._metrics.update()
            self._last_update = now

        registry = CollectorRegistry(auto_describe=True)

        Gauge("ejabberd_node_s2s_in", "count of incoming server-to-server connection", registry=registry).set(
            self._metrics.get_s2s_in()
        )
        Gauge("ejabberd_node_s2s_out", "count of outgoing server-to-server connection", registry=registry).set(
            self._metrics.get_s2s_out()
        )

        nodename = self._metrics.nodename
        Gauge("ejabberd_node_uptime", "uptime of ejabberd service", ["node"], registry=registry).labels(nodename).set(
            self._metrics.get_uptime()
        )
        Gauge("ejabberd_node_proccess", "count of pejabber proccess", ["node"], registry=registry).labels(nodename).set(
            self._metrics.get_processes()
        )

        labelnames_vhost = ["vhost"]

        registered_vhosts = Gauge(
            "ejabberd_registered_vhosts", "count of user per vhost", labelnames_vhost, registry=registry
        )
        muc = Gauge("ejabberd_muc", "count of muc's per vhost", labelnames_vhost, registry=registry)

        online_vhost_node = Gauge(
            "ejabberd_online_vhost_node", "count of client connections", ["vhost", "node"], registry=registry
        )

        online_status = Gauge(
            "ejabberd_online_status", "count of client connections", ["vhost", "node", "status"], registry=registry
        )
        online_connection = Gauge(
            "ejabberd_online_connection",
            "count of client connections",
            ["vhost", "node", "connection"],
            registry=registry,
        )
        online_client = Gauge(
            "ejabberd_online_client", "count of client software", ["vhost", "node", "client"], registry=registry
        )
        online_ipversion = Gauge(
            "ejabberd_online_ipversion", "count of client software", ["vhost", "node", "ipversion"], registry=registry
        )
        online_client_ipversion = Gauge(
            "ejabberd_online_client_ipversion",
            "count of client software",
            ["vhost", "node", "client", "ipversion"],
            registry=registry,
        )

        for host in self._metrics.get_vhosts():
            labels_vhost = host

            registered_vhosts.labels(labels_vhost).set(self._metrics.get_registered(host))
            muc.labels(labels_vhost).set(self._metrics.get_muc(host))

            for k, v in self._metrics.get_online_by_node(vhost=host).items():
                online_vhost_node.labels(host, k).set(v)

            for node in self._metrics.get_nodes():
                for k, v in self._metrics.get_online_by_status(node=node, vhost=host).items():
                    online_status.labels(host, node, k).set(v)
                for k, v in self._metrics.get_online_by_connection(node=node, vhost=host).items():
                    online_connection.labels(host, node, k).set(v)
                for k, v in self._metrics.get_online_by_client(node=node, vhost=host).items():
                    online_client.labels(host, node, k).set(v)
                for k, v in self._metrics.get_online_by_ipversion(node=node, vhost=host).items():
                    online_ipversion.labels(host, node, k).set(v)
                for client, data in self._metrics.get_online_client_by_ipversion(node=node, vhost=host).items():
                    for k, v in data.items():
                        online_client_ipversion.labels(host, node, client, str(k)).set(v)

        return registry

    def listen(self, addr=("127.0.0.1", 8080)):
        if "::" in addr[0]:
            HTTPServer.address_family = AF_INET6
        server = ThreadingHTTPServer(addr, DynamicMetricsHandler.factory(self.handler))
        server.serve_forever()


if __name__ == "__main__":
    # load config
    config = Config()
    if config.get("debug", default=False):
        logging.getLogger().setLevel(logging.DEBUG)

    # credentials and parameters
    url = config.get("url", default="http://[::1]:5280/api")
    login = config.get("login", default=None)
    api = config.get("api", default="rest")

    # config prometheus
    prom_addr = config.get("prometheus_address", default="127.0.0.1")
    prom_port = config.get("prometheus_port", default=8080)

    metrics = EjabberdMetrics(url, login, api)
    prom = Prometheus(metrics)
    prom.ttl = config.get("prometheus_cache_ttl", default=10)
    prom.listen((prom_addr, prom_port))