quay/util/saas/metricqueue.py

import datetime
import logging
import time

from functools import wraps
from Queue import Queue, Full

from flask import g, request


logger = logging.getLogger(__name__)

class MetricQueue(object):
  def __init__(self):
    self._queue = None

  def enable(self, maxsize=10000):
    self._queue = Queue(maxsize)

  def put(self, name, value, **kwargs):
    if self._queue is None:
      logger.debug('No metric queue %s %s %s', name, value, kwargs)
      return

    try:
      kwargs.setdefault('timestamp', datetime.datetime.now())
      kwargs.setdefault('dimensions', {})
      self._queue.put_nowait((name, value, kwargs))
    except Full:
      logger.error('Metric queue full')

  def get(self):
    return self._queue.get()

  def get_nowait(self):
    return self._queue.get_nowait()

def time_blueprint(bp, metric_queue):
  bp.before_request(time_before_request)
  bp.after_request(time_after_request(bp.name, metric_queue))

def time_before_request():
  g._request_start_time = time.time()

def time_after_request(name, metric_queue):
  def f(r):
    start = getattr(g, '_request_start_time', None)
    if start is None:
      return r

    dur = time.time() - start
    dims = {'endpoint': request.endpoint}

    metric_queue.put('ResponseTime', dur, dimensions=dims, unit='Seconds')
    metric_queue.put('ResponseCode', r.status_code, dimensions=dims)

    if r.status_code >= 500:
      metric_queue.put('5XXResponse', 1, dimensions={'name': name})
    elif r.status_code < 200 or r.status_code >= 300:
      metric_queue.put('Non200Response', 1, dimensions={'name': name})

    return r
  return f

def time_decorator(name, metric_queue):
  after = time_after_request(name, metric_queue)
  def decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
      time_before_request()
      rv = func(*args, **kwargs)
      after(rv)
      return rv
    return wrapper
  return decorator