quay/util/metrics/prometheus.py

import logging
import time
import threading

from flask import g, request
from prometheus_client import push_to_gateway, REGISTRY, Histogram


logger = logging.getLogger(__name__)


request_duration = Histogram('quay_request_duration_seconds',
                             'seconds taken to process a request',
                             labelnames=['method', 'endpoint', 'status'],
                             buckets=[.01, .025, .05, .1, .25, .5, 1.0, 2.5, 5.0])


PROMETHEUS_PUSH_INTERVAL_SECONDS = 30
ONE_DAY_IN_SECONDS = 60 * 60 * 24


class PrometheusPlugin(object):
  """ Application plugin for reporting metrics to Prometheus. """
  def __init__(self, app=None):
    self.app = app
    if app is not None:
      self.state = self.init_app(app)
    else:
      self.state = None

  def init_app(self, app):
    pusher = ThreadPusher(app)
    pusher.start()

    # register extension with app
    app.extensions = getattr(app, 'extensions', {})
    app.extensions['prometheus'] = pusher
    return pusher

  def __getattr__(self, name):
    return getattr(self.state, name, None)


class ThreadPusher(threading.Thread):
  def __init__(self, app):
    super(ThreadPusher, self).__init__()
    self.daemon = True
    self._app = app

  def run(self):
    agg_url = self._app.config.get('PROMETHEUS_PUSHGATEWAY_URL')
    while True:
      if agg_url is None:
        # Practically disable this worker, if there is no pushgateway.
        time.sleep(ONE_DAY_IN_SECONDS)
        continue

      time.sleep(PROMETHEUS_PUSH_INTERVAL_SECONDS)
      push_to_gateway(agg_url, job=self._app.config.get('PROMETHEUS_NAMESPACE', 'quay'),
                      registry=REGISTRY)


def timed_blueprint(bp):
  """
  Decorates a blueprint to have its request duration tracked by Prometheus.
  """
  def _time_before_request():
    g._request_start_time = time.time()
  bp.before_request(_time_before_request)


  def _time_after_request():
    def f(r):
      start = getattr(g, '_request_start_time', None)
      if start is None:
        return r
      dur = time.time() - start
      request_duration.labels(request.method, request.endpoint, r.status_code).observe(dur)
      return r
    return f
  bp.after_request(_time_after_request())
  return bp