add prometheus endpoint to analysis server

pyinfra/server/server.py

@@ -3,6 +3,7 @@ from typing import Dict
 
 from flask import Flask, jsonify, request
 from funcy import merge
+from prometheus_client import generate_latest, Summary
 
 from pyinfra.server.buffering.stream import FlatStreamBuffer
 from pyinfra.server.stream.queued_stream_function import QueuedStreamFunction
@@ -54,6 +55,10 @@ def __set_up_processing_server_impl(operation2function: Dict[str, QueuedStreamFu
     pickup_operation2processor = {pickup_suffix(op): prc for op, prc in operation2processor.items()}
     operation2processor = merge(submit_operation2processor, pickup_operation2processor)
 
+    pickup_operation2metric = {
+        op: Summary(f"redactmanager_{op}_seconds", f"Time spent on {op}.") for op in operation2processor
+    }
+
     @app.route("/ready", methods=["GET"])
     def ready():
         resp = jsonify("OK")
@@ -66,12 +71,17 @@ def __set_up_processing_server_impl(operation2function: Dict[str, QueuedStreamFu
         resp.status_code = 200
         return resp
 
+    @app.route("/prometheus", methods=["GET"])
+    def prometheus():
+        return generate_latest()
+
     @app.route("/<operation>", methods=["POST", "PATCH"])
     def submit(operation):
         return operation2processor[operation].push(request)
 
     @app.route("/<operation>", methods=["GET"])
     def pickup(operation):
-        return operation2processor[operation].pop()
+        with pickup_operation2metric[operation.replace("_pickup", "")].time():
+            return operation2processor[operation].pop()
 
     return app

requirements.txt

@@ -17,3 +17,4 @@ PyMuPDF==1.19.6
 more-itertools==8.12.0
 numpy==1.22.3
 Pillow==9.0.1
+prometheus-client==0.13.1