wsgi middleware for basic Prometheus monitoring instrumentation

middleware.wsgi

"""WSGI middleware for basic Prometheus monitoring instrumentation."""

import prometheus_client
import werkzeug.wrappers

Counter = prometheus_client.Counter
Histogram = prometheus_client.Histogram
INF = float('inf')


def PowersOf(logbase, count, lower=0, include_zero=True):
    """Returns a list of count powers of logbase (from logbase**lower)."""
    if not include_zero:
        return [logbase ** i for i in range(lower, count+lower)] + [INF]
    else:
        return [0] + [logbase ** i for i in range(lower, count+lower)] + [INF]


class Metrics(object):
    RequestCounter = Counter(
        'http_requests_total', 'Total number of HTTP requests.',
        ['method', 'scheme'])
    ResponseCounter = Counter(
        'http_responses_total', 'Total number of HTTP responses.', ['status'])
    LatencyHistogram = Histogram(
        'http_latency_seconds', 'Overall HTTP transaction latency.')
    RequestSizeHistogram = Histogram(
        'http_requests_body_bytes',
        'Breakdown of HTTP requests by content length.',
        buckets=PowersOf(2, 30))
    ResponseSizeHistogram = Histogram(
        'http_responses_body_bytes',
        'Breakdown of HTTP responses by content length.',
        buckets=PowersOf(2, 30))


class MetricsMiddleware(object):

    def __init__(self, app):
        self._app = app

    @Metrics.LatencyHistogram.time()
    def __call__(self, environ, start_response):
        request = werkzeug.wrappers.Request(environ)

        Metrics.RequestCounter.labels(request.method, request.scheme).inc()

        content_length = request.content_length
        if content_length is not None:
            Metrics.RequestSizeHistogram.observe(content_length)

        def counted_start_response(status, headers, exc_info=None):
            status_code = status.split(None, 1)[0]
            Metrics.ResponseCounter.labels(status_code).inc()
            try:
                content_length = int(dict(headers).get('Content-Length'))
                Metrics.ResponseSizeHistogram.observe(content_length)
            except (ValueError, TypeError):
                pass
            return start_response(status, headers, exc_info)

        response = self._app(environ, counted_start_response)
        return response

Hey! This looks like exactly like something I was just about to write myself. Have you considered putting this in a module?