1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74
|
import functools
from inspect import iscoroutinefunction
from time import perf_counter as time_now
def safe_wraps(wrapper, *args, **kwargs):
"""Safely wraps partial functions."""
while isinstance(wrapper, functools.partial):
wrapper = wrapper.func
return functools.wraps(wrapper, *args, **kwargs)
class Timer(object):
"""A context manager/decorator for statsd.timing()."""
def __init__(self, client, stat, rate=1):
self.client = client
self.stat = stat
self.rate = rate
self.ms = None
self._sent = False
self._start_time = None
def __call__(self, f):
"""Thread-safe timing function decorator."""
if iscoroutinefunction(f):
@safe_wraps(f)
async def _async_wrapped(*args, **kwargs):
start_time = time_now()
try:
return await f(*args, **kwargs)
finally:
elapsed_time_ms = 1000.0 * (time_now() - start_time)
self.client.timing(self.stat, elapsed_time_ms, self.rate)
return _async_wrapped
@safe_wraps(f)
def _wrapped(*args, **kwargs):
start_time = time_now()
try:
return f(*args, **kwargs)
finally:
elapsed_time_ms = 1000.0 * (time_now() - start_time)
self.client.timing(self.stat, elapsed_time_ms, self.rate)
return _wrapped
def __enter__(self):
return self.start()
def __exit__(self, typ, value, tb):
self.stop()
def start(self):
self.ms = None
self._sent = False
self._start_time = time_now()
return self
def stop(self, send=True):
if self._start_time is None:
raise RuntimeError('Timer has not started.')
dt = time_now() - self._start_time
self.ms = 1000.0 * dt # Convert to milliseconds.
if send:
self.send()
return self
def send(self):
if self.ms is None:
raise RuntimeError('No data recorded.')
if self._sent:
raise RuntimeError('Already sent data.')
self._sent = True
self.client.timing(self.stat, self.ms, self.rate)
|