Server-Timing header tools

server_timing.py

"""
Originally based on Jason Williams' (@jaswilli), timing work for GCSv5 API
"""

from timeit import default_timer as timer
import logging
from flask import g

log = logging.getLogger(__name__)


class TimingHeaderCollection(object):
    def __init__(self):
        self.headers = {}
        self.in_progress = set()

    def __getitem__(self, key):
        return self.headers[key]

    def __setitem__(self, key, val):
        self.headers[key] = val

    def force_in_progress(self):
        for x in self.in_progress:
            x.add_to_headers()

    def add_to_response(self, response):
        self.force_in_progress()
        # skip if empty
        if self.headers:
            log.debug(
                'Adding timing info to response: {}'.format(self.headers))
            response.headers['Server-Timing'] = (
                ','.join(['{}={}; "{}"'.format(k, v, name)
                          for (k, (name, v)) in self.headers.items()]))


class ServerTiming(object):
    def __init__(self, metric_name, nice_name=None):
        # get timing header collection, or, if absent, initialize it!
        self.timing_headers = g.get('timing_headers', None)
        if not self.timing_headers:
            g.timing_headers = TimingHeaderCollection()
            self.timing_headers = g.get('timing_headers')

        self.metric_name = metric_name
        self.metric_nice_name = nice_name or metric_name
        self.start_time = None

    def compute_elapsed_time(self):
        finish_time = timer()
        return round(((finish_time - self.start_time) * 1000), 2)

    def add_to_headers(self):
        elapsed_time_ms = self.compute_elapsed_time()
        log.debug('Setting timing val: {}, {}'
                  .format(self.metric_name, elapsed_time_ms))
        self.timing_headers[self.metric_name] = (self.metric_nice_name,
                                                 elapsed_time_ms)

    def __enter__(self):
        self.start_time = timer()
        self.timing_headers.in_progress.add(self)
        return self

    def __exit__(self, *args):
        self.add_to_headers()
        self.timing_headers.in_progress.remove(self)