ejabberd-tools/prometheus.py

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

from ejabberdrpc import EjabberdMetrics

import time
import threading
import socket
from http.server import BaseHTTPRequestHandler, HTTPServer
from socketserver import ThreadingMixIn

class _ThreadingSimpleServer(ThreadingMixIn, HTTPServer):
    """Thread per request HTTP server."""
    # Make worker threads "fire and forget". Beginning with Python 3.7 this
    # prevents a memory leak because ``ThreadingMixIn`` starts to gather all
    # non-daemon threads in a list in order to join on them at server close.
    # Enabling daemon threads virtually makes ``_ThreadingSimpleServer`` the
    # same as Python 3.7's ``ThreadingHTTPServer``.
    daemon_threads = True
    address_family = socket.AF_INET6


class Prometheus():
    def __init__(self, metrics):
        self._metrics = metrics
    
    def _parse_metric(self, name, value, tags=None):
        output = name
        if isinstance(tags, dict):
            output += "{"
            first = True
            for k, v in tags.items():
                if not first:
                    output += ', '
                else:
                    first = False
                output += k+'="'+v+'"'
            output += '}'
        return output + ' {}\n'.format(value)

    def _get_metrics(self):
        output = ""

        output += self._parse_metric("ejabberd_node_s2s_in", self._metrics.get_s2s_in())
        output += self._parse_metric("ejabberd_node_s2s_out", self._metrics.get_s2s_out())
        for host in self._metrics.get_vhosts():
            output += self._parse_metric("ejabberd_registered_vhosts", self._metrics.get_registered(host), {"vhost": host})
            muc = self._metrics.get_muc(host)
            if muc is not None:
                output += self._parse_metric("ejabberd_muc", muc, {"vhost": host})

            for k, v in self._metrics.get_online_by_node(vhost=host).items():
                output += self._parse_metric("ejabberd_online_vhost_node", v, {"vhost": host, "node": k})

            for node in self._metrics.get_nodes():
                for k, v in self._metrics.get_online_by_status(node=node, vhost=host).items():
                    output += self._parse_metric("ejabberd_online_status", v, {"vhost": host, "node": node, "status": k})
                for k, v in self._metrics.get_online_by_connection(node=node, vhost=host).items():
                    output += self._parse_metric("ejabberd_online_connection", v, {"vhost": host, "node": node, "connection": k})
                for k, v in self._metrics.get_online_by_client(node=node, vhost=host).items():
                    output += self._parse_metric("ejabberd_online_client", v, {"vhost": host, "node": node, "client": k})
                for k, v in self._metrics.get_online_by_ipversion(node=node, vhost=host).items():
                    output += self._parse_metric("ejabberd_online_ipversion", v, {"vhost": host, "node": node, "ipversion": str(k)})
                for client, data in self._metrics.get_online_client_by_ipversion(node=node,vhost=host).items():
                    for k, v in data.items():
                        output += self._parse_metric("ejabberd_online_client_ipversion", v, {"vhost": host, "node": node, "ipversion": str(k), "client": client})


        return output


    def listen(self, port, addr='::'):
        """Starts an HTTP server for prometheus metrics as a daemon thread"""
        class myHandler(BaseHTTPRequestHandler):
            def do_GET(r):
                r.send_response(200)
                r.send_header('Content-type', 'text/html')
                r.end_headers()
                result = self._get_metrics()
                r.wfile.write(result.encode('utf-8'))

        
        httpd = _ThreadingSimpleServer((addr, port), myHandler)
        t = threading.Thread(target=httpd.serve_forever)
        t.daemon = True
        t.start()

if __name__ == "__main__":
    import os
    import json

    # load config
    path = os.path.dirname(__file__)
    with open("/".join([path, "config.json"]), "r", encoding="utf-8") as f:
        config = json.load(f)


    url = config['url'] if "url" in config else "http://[::1]:5280/api"
    login = config['login'] if "login" in config else None
    api = config['api'] if "api" in config else "rest"

    prom_port = config['prometheus_port'] if "prometheus_port" in config else 8080
    prom_refresh = config['prometheus_refresh'] if "prometheus_refresh" in config else 10

    metrics = EjabberdMetrics(url, login, api)
    prom = Prometheus(metrics)
    prom.listen(prom_port)
    while True:
        metrics.update()
        time.sleep(prom_refresh)