Prometheus exporter that just sleeps a specified duration and returns that number as a metric

sleep_exporter.py

#!/usr/bin/env python3

import argparse
import logging
import os
from urllib.parse import parse_qs
import time
import wsgiref.simple_server
import wsgiref.util
from wsgiref.simple_server import ServerHandler

import prometheus_client
from prometheus_client.core import GaugeMetricFamily
from prometheus_client.exposition import ThreadingWSGIServer

class CustomWSGIRequestHandler(wsgiref.simple_server.WSGIRequestHandler):
    """Override the default log_message method to use unified logging"""

    def log_message(self, format, *args):
        logger = logging.getLogger("wsgiref")
        logger.debug(format % args)


class SleepExporter(prometheus_client.registry.Collector):
    def __init__(self, duration: str):
        self.duration = float(duration)

    def describe(self):
        return [GaugeMetricFamily("sleep", "Sleep duration in seconds", labels=["duration"])]

    def collect(self):
        logger.info(f"Sleeping {self.duration} seconds")
        time.sleep(self.duration)

        sleep = GaugeMetricFamily("sleep", "Sleep duration in seconds", labels=["duration"])
        sleep.add_metric([str(self.duration)], self.duration)
        yield sleep


def wsgi_app(environ, start_response):
    path = wsgiref.util.shift_path_info(environ)
    if path == "probe":
        try:
            qs = parse_qs(environ["QUERY_STRING"])
            duration = qs["duration"][0]
        except KeyError:
            start_response("400 Bad Request", [("Content-Type", "text/plain")])
            return [b"Duration parameter is missing\r\n"]

        registry = prometheus_client.registry.CollectorRegistry(auto_describe=False)
        registry.register(SleepExporter(duration))
        prometheus_app = prometheus_client.make_wsgi_app(registry)
        return prometheus_app(environ, start_response)

    else:
        start_response("404 Not Found", [("Content-Type", "text/plain")])
        return [b"Not Found. Use /probe endpoint.\r\n"]


def configure_logging(level: str) -> logging.Logger:
    """If running under systemd, format log messages for journald"""

    if "INVOCATION_ID" in os.environ:
        from systemd import journal # type: ignore
        logging.basicConfig(
            format="%(name)s:%(message)s",
            handlers=[journal.JournalHandler(SYSLOG_IDENTIFIER="sleep_exporter")],
            level=getattr(logging, level.upper())
        )
    else:
        logging.basicConfig(
            format="%(asctime)s [%(levelname)s] %(name)s:%(message)s",
            level=getattr(logging, level.upper()),
        )

    logger = logging.getLogger("sleep_exporter")
    return logger


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--web-listen-address", default="0.0.0.0")
    parser.add_argument("--web-listen-port", type=int, default=20006)
    parser.add_argument("--log-level", choices=["info", "debug", "error", "warning"], default="info", help="Set the logging level")
    args = parser.parse_args()

    logger = configure_logging(args.log_level)

    # Use HTTP 1.1 instead of 1.0 to reduce network traffic and CPU on Alloy
    # https://github.com/prometheus/client_python/issues/927
    ServerHandler.http_version = "1.1"

    httpd = wsgiref.simple_server.make_server(
        args.web_listen_address,
        args.web_listen_port,
        wsgi_app,
        ThreadingWSGIServer,
        handler_class=CustomWSGIRequestHandler,
    )
    logger.info("Starting sleep_exporter")
    logger.info(f"Listening on {args.web_listen_address}:{args.web_listen_port}")
    httpd.serve_forever()