add prometheus metrics for failed/successful connections

docker-compose.yml

@@ -6,3 +6,15 @@ services:
     ports:
       - "5000:5000"
 
+  prometheus:
+    image: prom/prometheus:latest
+    restart: always
+    volumes:
+      - ./prometheus.yml:/etc/prometheus/prometheus.yml
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+    ports:
+      - "9090:9090"
+
+
+

prometheus.yml

@@ -0,0 +1,7 @@
+global:
+  scrape_interval: 10s
+
+scrape_configs:
+  - job_name: 'smee2'
+    static_configs:
+      - targets: ['smee2:5000']

requirements.txt

@@ -1,2 +1,3 @@
 fastapi==0.136.3
-uvicorn==0.49.0
+uvicorn[standard]==0.49.0
+prometheus_client==0.21.1

server.py

@@ -1,7 +1,8 @@
-from fastapi import raiseExceptions
-from fastapi import FastAPI, WebSocket, Request, HTTPException
+from fastapi import FastAPI, WebSocket, Request, HTTPException, Response
+from prometheus_client import Counter, Gauge, generate_latest
 import uvicorn
 import logging
+import collections
 
 logging.basicConfig(
     format="%(asctime)s.%(msecs)03dZ %(levelname)s:%(name)s:%(message)s",
@@ -13,7 +14,19 @@
 
 app = FastAPI()
 
-clients = {}
+failed_connections = Counter (
+    "failed_connections",
+    "Number of failed connections to /tunnel",
+    ["subscription_id"]
+)
+
+successful_connections_per_subscription = Gauge (
+    "successful_connections",
+    "Number of successful connections to /tunnel",
+    ["subscription_id"]
+)
+
+clients = collections.defaultdict(list)
 subscribers = {}
 
 @app.post("/webhook/{subscription_id}")
@@ -27,13 +40,13 @@ async def webhook(subscription_id: str, request: Request):
         logging.info("Webhook received: %s", data)
 
         subscribers[subscription_id] = data
-        client = clients.get(subscription_id)
-
-        if client is not None:
-            await client.send_json(data) 
-            print("Data sent to websocket client")
+
+        for client in clients.get(subscription_id, []):
+            if client is not None:
+                await client.send_json(data)
+                print("Data sent to websocket client") 
         return {"message":"received"}  
-
+                   
     else:   
         print("Invalid endpoint, connection not accepted")
         return
@@ -42,15 +55,27 @@ async def webhook(subscription_id: str, request: Request):
 @app.websocket("/tunnel/{subscription_id}")
 async def websocket_endpoint(subscription_id: str, websocket: WebSocket):
     await websocket.accept()
-    clients[subscription_id] = websocket
+    clients[subscription_id].append(websocket)
     try:
         while True:
             data = await websocket.receive_text()
             await websocket.send_text("Message received")
+            successful_connections_per_subscription.labels(subscription_id).inc()
+
     except Exception as e:
-        clients.pop(subscription_id, None)
+        clients[subscription_id].remove(websocket)
+        if not clients[subscription_id]:
+            clients.pop(subscription_id, None)
+        failed_connections.labels(subscription_id).inc()
+
 
+@app.get("/metrics")
+def get_metrics():
+    return Response(
+        content=generate_latest(),
+        media_type="text/plain"
+    )
 
 if __name__ == "__main__":
-    uvicorn.run("server:app", host="0.0.0.0", port=5000) 
+    uvicorn.run(app, host="0.0.0.0", port=5000)