diff --git a/docker-compose.yml b/docker-compose.yml index a625591..63bf446 100644 --- a/docker-compose.yml +++ b/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" + + + diff --git a/prometheus.yml b/prometheus.yml new file mode 100644 index 0000000..a243672 --- /dev/null +++ b/prometheus.yml @@ -0,0 +1,7 @@ +global: + scrape_interval: 10s + +scrape_configs: + - job_name: 'smee2' + static_configs: + - targets: ['smee2:5000'] diff --git a/requirements.txt b/requirements.txt index 36f0877..2489a5e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ fastapi==0.136.3 -uvicorn==0.49.0 +uvicorn[standard]==0.49.0 +prometheus_client==0.21.1 \ No newline at end of file diff --git a/server.py b/server.py index 9928eca..1f3c764 100644 --- a/server.py +++ b/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) \ No newline at end of file