rr endpoint

balancer/util.py

@@ -5,22 +5,28 @@ import requests
 import json
 
 
-def proxy_datacenters(serverpath: str, username, method="GET", balancer_fun=rr_get):
+def proxy_datacenters(
+    serverpath: str, username, method="GET", body="", balancer_fun=rr_get
+):
+    if "datacenter" in body:
+        server = body["datacenter"]
+    else:
         server = balancer_fun()
     token = get_datacenter_token(username, server)
     url = f"{server}/{serverpath}"
     t_resp = requests.request(
+        json=body,
         method=method,
         url=url,
         allow_redirects=False,
         verify=False,
         headers={"Authorization": token},
     )
-    if t_resp.status_code / 100 != 2:
+    if int(t_resp.status_code / 100) != 2:
         raise HTTPException(
             status_code=t_resp.status_code, detail="Remote server error"
         )
-    response = ORJSONResponse(
-        json.loads(t_resp.content), status_code=t_resp.status_code
-    )
+    obj = json.loads(t_resp.content)
+    obj[0]["datacenter"] = server
+    response = ORJSONResponse(obj[0], status_code=t_resp.status_code)
     return response

core/models.py

@@ -19,6 +19,7 @@ class UserLoginSchema(BaseModel):
     class Config:
         schema_extra = {"example": {"username": "user", "password": "weakpassword"}}
 
+
 class DataCenterSchema(BaseModel):
     name: str = Field(min_length=3, max_length=50)
     url: str

main.py

@@ -2,7 +2,13 @@ from fastapi import FastAPI, Response, Body, Depends
 import os
 from typing import List
 from balancer.util import proxy_datacenters
-from sredis.sredis import check_user, create_puser, add_datacenter, set_token, create_superpuser
+from sredis.sredis import (
+    check_user,
+    create_puser,
+    add_datacenter,
+    set_token,
+    create_superpuser,
+)
 import logging
 import requests
 from core.models import User, Token, UserLoginSchema, DataCenterSchema
@@ -22,7 +28,7 @@ add_datacenter("https://kappa1.fured.cloud.bme.hu")
 add_datacenter("https://kappa2.fured.cloud.bme.hu")
 add_datacenter("https://kappa3.fured.cloud.bme.hu")
 create_superpuser(
-    User(username='admin', email='example@domain.com', password=os.getenv('ADMIN_PASS'))
+    User(username="admin", email="example@domain.com", password=os.getenv("ADMIN_PASS"))
 )
 
 
@@ -40,13 +46,22 @@ async def user_login(user: UserLoginSchema = Body(...)):
 
 
 @app.get("/lb/{server_path:path}")
-def proxy_get(server_path: str = "/", username=Depends(get_current_user)):
-    return proxy_datacenters(server_path, username)
+def proxy_get(server_path: str = "/", username=Depends(get_current_user), body=Body()):
+    return proxy_datacenters(server_path, username, body=body)
+
+
+@app.post("/lb/rr/{server_path:path}")
+def proxy_post_rr(
+    server_path: str = "/", username=Depends(get_current_user), body=Body()
+):
+    return proxy_datacenters(server_path, username, body=body, method="POST")
 
 
-@app.post("/lb/{server_path:path}")
-def proxy_post(server_path: str = "/", username=Depends(get_current_user)):
-    return proxy_datacenters(server_path, username, method="POST")
+@app.post("/lb/ff/{server_path:path}")
+def proxy_post_ff(
+    server_path: str = "/", username=Depends(get_current_user), body=Body()
+):
+    return proxy_datacenters(server_path, username, body=body, method="POST")
 
 
 @app.post("/add_datacenter/")

sredis/models.py

@@ -8,6 +8,7 @@ class PUser(HashModel):
     password: str
     admin: int = 0
 
+
 class DataCenter(HashModel):
     name: str = Field(index=True)
     url: str = Field(index=True)

sredis/sredis.py

@@ -20,6 +20,7 @@ def add_datacenter(datacenter: str):
     r.incr("datacenters_cnt")
     r.set("roundrobin_cnt", 1)
 
+
 def rr_get():
     cnt = int(r.get("datacenters_cnt"))
     rr = int(r.get("roundrobin_cnt"))
@@ -29,11 +30,13 @@ def rr_get():
         r.incr("roundrobin_cnt")
     return str(r.hget("datacenters_hash", rr))
 
+
 def wr_get(centers: DataCenterSchema):
     cnt = int(r.get("datacenters_cnt"))
-    indexis = range(1, cnt+1)
+    indexis = range(1, cnt + 1)
     return random.choice(indexis, weights=())[0]
 
+
 def set_token(username: str, datacenter: str, token: str):
     print(f"tokens:{username}" + datacenter)
     r.hset(f"tokens:{username}", datacenter, token)
@@ -64,7 +67,10 @@ def create_superpuser(user: User):
     s = PUser.find(PUser.username == user.username).all()
     if not s:
         puser = PUser(
-            username=user.username, email=user.email, password=hash_pass(user.password), admin=1
+            username=user.username,
+            email=user.email,
+            password=hash_pass(user.password),
+            admin=1,
         )
         puser.save()
         return True

stresstest/main.py

 import json
 from locust import TaskSet, task, HttpUser
 
+
 class PerformanceTests(TaskSet):
     def on_start(self):
-        resp = self.client.post("/user/login", json={
-            "username": "karsa",
-            "password": "12345678"
-        })
-        self.client.headers = {'Authorization': 'Bearer ' + json.loads(resp._content)['access_token']} 
+        resp = self.client.post(
+            "/user/login", json={"username": "karsa", "password": "12345678"}
+        )
+        self.client.headers = {
+            "Authorization": "Bearer " + json.loads(resp._content)["access_token"]
+        }
 
     @task(1)
     def testFastApi(self):