quay/util/saas/metricqueue.py

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 = Queue(10000)

  def put(self, *args, **kwargs):
    try:
      self._queue.put_nowait((args, kwargs))
    except Full:
      logger.error('Metric queue full')

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

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._start = time.time()

def time_after_request(name, metric_queue):
  def f(r):
    start = getattr(g, '_start', None)
    if start is None:
      return r
    dur = time.time() - start
    dims = dimensions={'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 < 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