update rr method

balancer/util.py

@@ -33,6 +33,8 @@ def proxy_datacenters(
         raise HTTPException(
             status_code=t_resp.status_code, detail="Remote server error"
         )
+    if not t_resp.content:
+        return ORJSONResponse({"status": "ok"}, status_code=t_resp.status_code)
     obj = json.loads(t_resp.content)
     if type(obj) is list:
         for o in obj:

main.py

@@ -8,6 +8,8 @@ from sredis.sredis import (
     add_datacenter,
     set_token,
     create_superpuser,
+    get_rtokens, 
+    reset
 )
 import logging
 import requests
@@ -31,6 +33,8 @@ create_superpuser(
     User(username="admin", email="example@domain.com", password=os.getenv("ADMIN_PASS"))
 )
 
+reset()
+
 
 @app.post("/user/signup", tags=["user"])
 async def create_user(user: User = Body(...)):
@@ -84,3 +88,7 @@ def set_tokens(tokens: List[Token] = None, username=Depends(get_current_user)):
     for token in tokens:
         set_token(username, str(token.datacenter), str(token.token))
     return tokens
+
+@app.get("/tokens/")
+def get_tokens(username=Depends(get_current_user)):
+    return get_rtokens(username=username)

sredis/sredis.py

@@ -8,27 +8,30 @@ from fastapi import HTTPException
 import random
 
 r = redis.Redis(host="localhost", port=6379, db=0, decode_responses=True)
-r.set("datacenters_cnt", 1)
-all_keys = list(r.hgetall("datacenters_hash").keys())
-if all_keys:
-    r.hdel("datacenters_hash", *all_keys)
 
+def reset():
+    print("Reset OK")
+    r.set("datacenters_cnt", 1)
+    all_keys = r.hgetall("datacenters_hash")
+    if all_keys.keys():
+        r.hdel("datacenters_hash",*list(all_keys.keys()))
+
+    r.delete("rrlist")
+    pipe = r.pipeline()
+    for dc in all_keys.values():
+        pipe.rpush("rrlist", dc)
+    pipe.execute()
 
 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)
 
 def rr_get():
-    cnt = int(r.get("datacenters_cnt"))
-    rr = int(r.get("roundrobin_cnt"))
-    if rr + 1 >= cnt:
-        r.set("roundrobin_cnt", 1)
-    else:
-        r.incr("roundrobin_cnt")
-    return str(r.hget("datacenters_hash", rr))
+    rr = r.rpoplpush("rrlist", "rrlist")
+    return rr
 
 
 def wr_get(centers: DataCenterSchema):
@@ -40,6 +43,18 @@ def wr_get(centers: DataCenterSchema):
 def set_token(username: str, datacenter: str, token: str):
     r.hset(f"tokens:{username}", datacenter, token)
 
+def get_rtoken(username: str, datacenter: str):
+    return r.hget(f"tokens:{username}", datacenter)
+
+def get_rtokens(username: str):
+    center_ids = list(r.hgetall("datacenters_hash").keys())
+    tokens = []  
+    for id in center_ids:
+        token = { } 
+        token["datacenter"] = r.hget("datacenters_hash", id)
+        token["token"] = get_rtoken(username=username, datacenter=token["datacenter"])
+        tokens.insert(0, token)
+    return tokens
 
 def get_datacenter_token(username: str, datacenter: str):
     return str(r.hget(f"tokens:{username}", datacenter))

stresstest/main.py

@@ -5,7 +5,7 @@ from locust import TaskSet, task, HttpUser
 class PerformanceTests(TaskSet):
     def on_start(self):
         resp = self.client.post(
-            "/user/login", json={"username": "karsa", "password": "12345678"}
+            "/user/login", json={"username": "admin", "password": "12345678"}
         )
         self.client.headers = {
             "Authorization": "Bearer " + json.loads(resp._content)["access_token"]
@@ -15,6 +15,14 @@ class PerformanceTests(TaskSet):
     def testFastApi(self):
         self.client.get("/lb/dashboard/acpi/vm/1/")
 
+    @task(2)
+    def testRRLoadBalancerLease(self):
+        self.client.get("/lb/dashboard/acpi/lease")
+
+    @task(2)
+    def testRRLoadBalancerVMAll(self):
+        self.client.get("/lb/dashboard/acpi/vm")    
+
 
 class WebsiteUser(HttpUser):
     tasks = [PerformanceTests]