40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from functools import lru_cache
|
|
|
|
from funcy import identity
|
|
from prometheus_client import CollectorRegistry, Summary
|
|
|
|
from pyinfra.server.operation_dispatcher import OperationDispatcher
|
|
|
|
|
|
class OperationDispatcherMonitoringDecorator:
|
|
def __init__(self, operation_dispatcher: OperationDispatcher, naming_policy=identity):
|
|
self.operation_dispatcher = operation_dispatcher
|
|
self.operation2metric = {}
|
|
self.naming_policy = naming_policy
|
|
|
|
@property
|
|
@lru_cache(maxsize=None)
|
|
def registry(self):
|
|
return CollectorRegistry(auto_describe=True)
|
|
|
|
def make_summary_instance(self, op: str):
|
|
return Summary(f"{self.naming_policy(op)}_seconds", f"Time spent on {op}.", registry=self.registry)
|
|
|
|
def submit(self, operation, request):
|
|
return self.operation_dispatcher.submit(operation, request)
|
|
|
|
def pickup(self, operation):
|
|
with self.get_monitor(operation):
|
|
return self.operation_dispatcher.pickup(operation)
|
|
|
|
def get_monitor(self, operation):
|
|
monitor = self.operation2metric.get(operation, None) or self.register_operation(operation)
|
|
return monitor.time()
|
|
|
|
def register_operation(self, operation):
|
|
summary = self.make_summary_instance(operation)
|
|
self.operation2metric[operation] = summary
|
|
return summary
|
|
|
|
|