status setup

balancer/util.py

@@ -42,4 +42,4 @@ def proxy_datacenters(
     else:
         obj["datacenter"] = server
     response = ORJSONResponse(obj, status_code=t_resp.status_code)
-    return response
+    return response
\ No newline at end of file

main.py

@@ -10,9 +10,11 @@ from sredis.sredis import (
     create_superpuser,
     get_rtokens,
     reset,
+    update_datacenter_status,
 )
 import logging
 import requests
+from fastapi_utils.tasks import repeat_every
 from core.models import User, Token, UserLoginSchema, DataCenterSchema
 from core.auth import signJWT
 from core.bearer import get_current_user, admin_user
@@ -95,3 +97,9 @@ def set_tokens(tokens: List[Token] = None, username=Depends(get_current_user)):
 @app.get("/tokens/")
 def get_tokens(username=Depends(get_current_user)):
     return get_rtokens(username=username)
+
+
+@app.on_event("startup")
+@repeat_every(seconds=300)
+def update_datacenter_healt():
+    update_datacenter_status()

pyproject.toml

@@ -20,6 +20,7 @@ pydantic = {extras = ["email"], version = "^1.10.6"}
 passlib = "^1.7.4"
 black = "^23.1.0"
 ruff = "^0.0.254"
+fastapi-utils = "^0.2.1"
 
 [tool.poe.tasks.start]
 shell = "poetry run uvicorn main:app --reload --port 6973 --host 0.0.0.0"

sredis/sredis.py

 import redis
+import requests
 from core.models import UserLoginSchema
 from .models import PUser
 from passlib.hash import pbkdf2_sha256
@@ -7,6 +8,14 @@ from core.auth import hash_pass
 from fastapi import HTTPException
 import random
 
+def get_status(datacenter_url: str):
+    resp = requests.request(
+        method="GET",
+        url=datacenter_url,
+        verify=False
+    )
+    return int(resp.status_code / 100) == 2
+
 r = redis.Redis(host="localhost", port=6379, db=0, decode_responses=True)
 
 
@@ -24,12 +33,24 @@ def reset():
     pipe.execute()
 
 
+def update_status(datacenter):
+    status = "OK" if get_status(datacenter_url=datacenter) is True else "Wrong"
+    r.hset("datacenter_status", datacenter, status)
+
+
+def update_datacenter_status():
+    all_keys = r.hgetall("datacenters_hash")
+    for dc in all_keys.values():
+        update_status(datacenter=dc)
+
+
 def add_datacenter(datacenter: str):
     cnt = int(r.get("datacenters_cnt"))
     r.hset("datacenters_hash", cnt, datacenter)
     r.incr("datacenters_cnt")
     r.set("roundrobin_cnt", 1)
     r.rpush("rrlist", datacenter)
+    update_status(datacenter)
 
 
 def rr_get():