add health check endpoint
This commit is contained in:
parent
38eef0e33e
commit
b7a78ba140
53
ipserv.py
53
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)
|
||||
|
|
Reference in a new issue