From 9a9e2888d8945a6e8396a156176466607b3c7c31 Mon Sep 17 00:00:00 2001
From: Aravinth Manivannan <realaravinth@batsense.net>
Date: Sun, 17 Dec 2023 19:19:56 +0530
Subject: [PATCH] feat: benchmark pipeline and individual write endpoints

---
 bench/locustfile.py | 95 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 bench/locustfile.py

diff --git a/bench/locustfile.py b/bench/locustfile.py
new file mode 100644
index 0000000..1645bb0
--- /dev/null
+++ b/bench/locustfile.py
@@ -0,0 +1,95 @@
+from pprint import pprint
+from locust import FastHttpUser, between, task
+
+password = "fooobarasdfasdf"
+username = "realaravinth"
+
+
+class Unprotected(FastHttpUser):
+    #    wait_time = between(5, 15)
+    peers = [
+            "http://localhost:9001",
+            "http://localhost:9002",
+            "http://localhost:9003",
+            ]
+    leader = "http://localhost:9001"
+    host = leader
+    captcha_id="locust"
+
+    pipeline_vote = []
+    for _ in range(0,1000):
+        pipeline_vote.append({"AddVisitor": captcha_id})
+
+#    def on_start(self):
+#        resp  = self.client.get(f"{self.leader}/metrics")
+#        data = resp.json()
+#        leader = data["Ok"]["membership_config"]["log_id"]["leader_id"]["node_id"]
+#        self.leader = self.peers[leader - 1]
+#        self.host = self.leader
+#        print(f"Leader: {self.host}")
+#        self.add_captcha(captcha_id="locust")
+
+
+
+    def write(self, data):
+        resp = self.client.post(f"{self.host}/write", json=data)
+        print(f"RPC Status: {resp.status_code}")
+        resp = resp.json()
+        if "Err" in resp:
+            leader = resp["Err"]["APIError"]["ForwardToLeader"]["leader_node"]["addr"]
+            print(f"Forwarding write to leader {leader}")
+            return write(leader, data)
+        return resp["Ok"]["data"]
+
+    def pipeline_write(self, data):
+        resp = self.client.post(f"{self.host}/pipeline/write", json=data)
+#        print(f"RPC Status: {resp.status_code}")
+        resp = resp.json()
+        if "Err" in resp:
+            leader = resp["Err"]["APIError"]["ForwardToLeader"]["leader_node"]["addr"]
+            print(f"Forwarding write to leader {leader}")
+            return write(leader, data)
+        return resp
+
+
+
+
+    def add_vote(self, captcha_id: str):
+        resp = self.write(data={"AddVisitor": captcha_id})
+        pprint(resp)
+
+    def add_vote_pipeline(self, captcha_id: str):
+        resp = self.pipeline_write(data=self.pipeline_vote)
+#        pprint(resp)
+
+    def add_captcha(self, captcha_id: str):
+        params = {
+            "AddCaptcha": {
+                "id": captcha_id,
+                "mcaptcha": {
+                    "visitor_threshold": 0,
+                    "defense": {
+                        "levels": [
+                            {"visitor_threshold": 50, "difficulty_factor": 500},
+                            {"visitor_threshold": 5000, "difficulty_factor": 50000},
+                        ],
+                        "current_visitor_threshold": 0,
+                    },
+                    "duration": 30,
+                },
+            }
+        }
+        resp = self.write(data=params)
+        pprint(f"Captcha added {captcha_id}: {resp}")
+
+
+    @task
+    def unprotected(self):
+        self.add_vote_pipeline(captcha_id=self.captcha_id)
+        ##self.add_vote(captcha_id="locust")
+#        data = {
+#            "username": username,
+#            "password": username,
+#            "confirm_password": username,
+#        }
+#        self.client.post("/unprotected", data=data)