refactored server setup code: factored out and decoupled operation registry and prometheus summary registry

pyinfra/server/server.py

@@ -1,8 +1,9 @@
 from functools import singledispatch, lru_cache
+from itertools import starmap, tee
 from typing import Dict, Callable, Union
 
 from flask import Flask, jsonify, request
-from funcy import merge
+from funcy import zipdict, juxt, cat
 from prometheus_client import generate_latest, Summary, CollectorRegistry
 
 from pyinfra.config import CONFIG
@@ -50,56 +51,85 @@ def __stream_fn_to_processing_server(operation2stream_fn: dict, buffer_size):
     return __set_up_processing_server(operation2stream_fn)
 
 
-def __set_up_processing_server(operation2function: Dict[str, QueuedStreamFunction]):
+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 get_registry():
+    def registry(self):
         return CollectorRegistry(auto_describe=True)
 
-    def make_summary_instance(op: str):
+    def make_summary_instance(self, op: str):
         op = op.replace("_pickup", "") if op != "pickup" else "default"
         service_name = CONFIG.service.name
-        return Summary(f"redactmanager_{service_name}_{op}_seconds", f"Time spent on {op}.", registry=get_registry())
+        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__)
 
-    operation2processor = {
-        op: LazyRestProcessor(fn, **build_endpoint_suffixes(op)) for op, fn in operation2function.items()
-    }
+    dispatcher = OperationDispatcherMonitoringDecorator(OperationDispatcher(operation2function))
 
-    submit_operation2processor = {submit_suffix(op): prc for op, prc in operation2processor.items()}
-    pickup_operation2processor = {pickup_suffix(op): prc for op, prc in operation2processor.items()}
-    operation2processor = merge(submit_operation2processor, pickup_operation2processor)
-
-    operation2metric = {op: make_summary_instance(op) for op in pickup_operation2processor}
+    def ok():
+        resp = jsonify("OK")
+        resp.status_code = 200
+        return resp
 
     @app.route("/ready", methods=["GET"])
     def ready():
-        resp = jsonify("OK")
-        resp.status_code = 200
-        return resp
+        return ok()
 
     @app.route("/health", methods=["GET"])
     def healthy():
-        resp = jsonify("OK")
-        resp.status_code = 200
-        return resp
+        return ok()
 
     @app.route("/prometheus", methods=["GET"])
     def prometheus():
-        return generate_latest(registry=get_registry())
+        return generate_latest(registry=dispatcher.registry)
 
     @app.route("/<operation>", methods=["POST", "PATCH"])
     def submit(operation):
-        return operation2processor[operation].push(request)
+        return dispatcher.push(operation, request)
 
     @app.route("/", methods=["POST", "PATCH"])
     def submit_default():
-        return operation2processor[""].push(request)
+        return dispatcher.push("", request)
 
     @app.route("/<operation>", methods=["GET"])
     def pickup(operation):
-        with operation2metric[operation].time():
-            return operation2processor[operation].pop()
+        return dispatcher.pop(operation)
 
     return app