add health check endpoint

ipserv.py

@@ -1,7 +1,7 @@
-import requests
 import json
 import logging
 import time
+import os
 import aiohttp
 import random
 from fastapi import FastAPI, Header, Request, Query, HTTPException
@@ -11,6 +11,23 @@ from contextlib import asynccontextmanager
 logging.basicConfig(level=logging.INFO)
 
 
+class Manager:
+    def __init__(self, api: FastAPI):
+        self.app = api
+        self.waiting = 0
+
+    def __enter__(self) -> "Manager":
+        self.waiting += 1
+        if self.waiting >= 2048:
+            logging.critical("TCP pool full! %d/2048. Requests are now being backlogged.")
+        elif self.waiting > 1024:
+            logging.warning("TCP pool half full! %d/2048.")
+        return self
+
+    def __exit__(self, *args):
+        self.waiting -= 1
+
+
 @asynccontextmanager
 async def lifespan(_app: FastAPI):
     async with aiohttp.ClientSession(
@@ -24,6 +41,15 @@ async def lifespan(_app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 app.state.cache = {}
+app.state.meta = Manager(app)
+
+
+@app.middleware("http")
+async def http_middleware(request: Request, call_next):
+    if app.state.meta.waiting >= 2512:
+        return JSONResponse({"error": os.urandom(512).decode("latin-1", "replace")}, status_code=503)
+    with app.state.meta:
+        return await call_next(request)
 
 
 async def make_request(ip: str, headers: dict[str, str]) -> dict | HTTPException:
@@ -45,6 +71,7 @@ async def make_request(ip: str, headers: dict[str, str]) -> dict | HTTPException
         return HTTPException(500, "Failed to get upstream data.")
     return data
 
+
 @app.get("/")
 async def ip(
     request: Request,
@@ -125,3 +152,27 @@ async def im_feeling_lucky(req: Request):
 @app.get("/raw")
 def get_raw(req: Request):
     return PlainTextResponse(req.client.host)
+
+
+@app.get("/health")
+def get_health():
+    detail = {"issues": []}
+    if app.state.meta.waiting >= 2048:
+        detail["status"] = "critical"
+        detail["issues"].append("(C) Connection pool full.")
+    elif app.state.meta.waiting >= 1024:
+        detail["status"] = "warning"
+        detail["issues"].append("(W) TCP pool half full.")
+    else:
+        detail["status"] = "ok"
+
+    try:
+        t = time.perf_counter()
+        async with app.state.session.get("/") as response:
+            await response.text()
+        detail["latency"] = time.perf_counter() - t
+    except Exception as e:
+        detail["issues"].append(f"(E) Failed to check upstream: {e}")
+        detail["status"] = "critical"
+    
+    return JSONResponse(detail)

tox.ini

@@ -0,0 +1,2 @@
+[flake8]
+max-line-length = 128