diff --git a/pyinfra/server/monitoring.py b/pyinfra/server/monitoring.py new file mode 100644 index 0000000..1122461 --- /dev/null +++ b/pyinfra/server/monitoring.py @@ -0,0 +1,37 @@ +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 push(self, operation, request): + return self.operation_dispatcher.push(operation, request) + + def register_operation(self, operation): + summary = self.make_summary_instance(operation) + self.operation2metric[operation] = summary + return summary + + def get_monitor(self, operation): + monitor = self.operation2metric.get(operation, None) or self.register_operation(operation) + return monitor.time() + + def pop(self, operation): + with self.get_monitor(operation): + return self.operation_dispatcher.pop(operation) diff --git a/pyinfra/server/operation_dispatcher.py b/pyinfra/server/operation_dispatcher.py new file mode 100644 index 0000000..90f175d --- /dev/null +++ b/pyinfra/server/operation_dispatcher.py @@ -0,0 +1,34 @@ +from itertools import starmap, tee +from typing import Dict + +from funcy import juxt, zipdict, cat + +from pyinfra.server.stream.queued_stream_function import QueuedStreamFunction +from pyinfra.server.stream.rest import LazyRestProcessor + + +class OperationDispatcher: + + def __init__(self, operation2function: Dict[str, QueuedStreamFunction]): + submit_suffixes, pickup_suffixes = zip(*map(juxt(submit_suffix, pickup_suffix), operation2function)) + processors = starmap(LazyRestProcessor, zip(operation2function.values(), submit_suffixes, pickup_suffixes)) + self.operation2processor = zipdict(submit_suffixes + pickup_suffixes, cat(tee(processors))) + + @classmethod + @property + def pop_suffix(cls): + return pickup_suffix("") + + def push(self, operation, request): + return self.operation2processor[operation].push(request) + + def pop(self, operation): + return self.operation2processor[operation].pop() + + +def submit_suffix(op: str): + return "" if not op else op + + +def pickup_suffix(op: str): + return "pickup" if not op else f"{op}_pickup" diff --git a/pyinfra/server/server.py b/pyinfra/server/server.py index 0d39926..5be345f 100644 --- a/pyinfra/server/server.py +++ b/pyinfra/server/server.py @@ -1,15 +1,14 @@ -from functools import singledispatch, lru_cache -from itertools import starmap, tee +from functools import singledispatch from typing import Dict, Callable, Union from flask import Flask, jsonify, request -from funcy import zipdict, juxt, cat -from prometheus_client import generate_latest, Summary, CollectorRegistry +from prometheus_client import generate_latest from pyinfra.config import CONFIG from pyinfra.server.buffering.stream import FlatStreamBuffer +from pyinfra.server.monitoring import OperationDispatcherMonitoringDecorator +from pyinfra.server.operation_dispatcher import OperationDispatcher from pyinfra.server.stream.queued_stream_function import QueuedStreamFunction -from pyinfra.server.stream.rest import LazyRestProcessor @singledispatch @@ -51,56 +50,13 @@ def __stream_fn_to_processing_server(operation2stream_fn: dict, buffer_size): return __set_up_processing_server(operation2stream_fn) -class OperationDispatcher: - def __init__(self, operation2function: Dict[str, QueuedStreamFunction]): - submit_suffixes, pickup_suffixes = zip(*map(juxt(submit_suffix, pickup_suffix), operation2function)) - processors = starmap(LazyRestProcessor, zip(operation2function.values(), submit_suffixes, pickup_suffixes)) - self.operation2processor = zipdict(submit_suffixes + pickup_suffixes, cat(tee(processors))) - - def push(self, operation, request): - return self.operation2processor[operation].push(request) - - def pop(self, operation): - return self.operation2processor[operation].pop() - - -class OperationDispatcherMonitoringDecorator: - def __init__(self, operation_dispatcher: OperationDispatcher, prefix="redactmanager"): - self.operation_dispatcher = operation_dispatcher - self.prefix = prefix - self.operation2metric = {} - - @property - @lru_cache(maxsize=None) - def registry(self): - return CollectorRegistry(auto_describe=True) - - def make_summary_instance(self, op: str): - op = op.replace("_pickup", "") if op != "pickup" else "default" - service_name = CONFIG.service.name - return Summary(f"{self.prefix}_{service_name}_{op}_seconds", f"Time spent on {op}.", registry=self.registry) - - def push(self, operation, request): - return self.operation_dispatcher.push(operation, request) - - def register_operation(self, operation): - summary = self.make_summary_instance(operation) - self.operation2metric[operation] = summary - return summary - - def get_monitor(self, operation): - monitor = self.operation2metric.get(operation, None) or self.register_operation(operation) - return monitor.time() - - def pop(self, operation): - with self.get_monitor(operation): - return self.operation_dispatcher.pop(operation) - - def __set_up_processing_server(operation2function: Dict[str, QueuedStreamFunction]): app = Flask(__name__) - dispatcher = OperationDispatcherMonitoringDecorator(OperationDispatcher(operation2function)) + dispatcher = OperationDispatcherMonitoringDecorator( + OperationDispatcher(operation2function), + naming_policy=naming_policy, + ) def ok(): resp = jsonify("OK") @@ -134,13 +90,11 @@ def __set_up_processing_server(operation2function: Dict[str, QueuedStreamFunctio return app -def build_endpoint_suffixes(op: str): - return {"submit_suffix": submit_suffix(op), "pickup_suffix": pickup_suffix(op)} +def naming_policy(op_name: str): + pop_suffix = OperationDispatcher.pop_suffix + prefix = f"redactmanager_{CONFIG.service.name}" + op_display_name = op_name.replace(f"_{pop_suffix}", "") if op_name != pop_suffix else "default" + complete_display_name = f"{prefix}_{op_display_name}" -def submit_suffix(op: str): - return "" if not op else op - - -def pickup_suffix(op: str): - return "pickup" if not op else f"{op}_pickup" + return complete_display_name