dcache/bench/locustfile.py

import json
import time
import grpc
from pprint import pprint
from locust import FastHttpUser, between, task

from dcache_py import dcache_pb2 as dcache
from dcache_py.dcache_pb2 import RaftRequest
from dcache_py.dcache_pb2_grpc import DcacheServiceStub

# 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)
#
import gevent
import adaptor
from locust import events, task

host = "localhost:9001"
captcha_id = "locust"


def add_captcha(stub: DcacheServiceStub, captcha_id: str):
    msg = dcache.AddCaptchaRequest(
        id=captcha_id,
        mcaptcha=dcache.MCaptcha(
            duration=30,
            defense=dcache.Defense(
                levels=[
                    dcache.Level(visitor_threshold=50, difficulty_factor=500),
                    dcache.Level(visitor_threshold=5000, difficulty_factor=50000),
                ]
            ),
        ),
    )
    resp = stub.AddCaptcha(msg)
    pprint(f"Captcha added {captcha_id}: {resp}")


with grpc.insecure_channel(host) as channel:
    stub = DcacheServiceStub(channel)
    add_captcha(stub=stub, captcha_id=captcha_id)


#pipeline_msgs = []
#for _ in range(0,1000):
#    pipeline_msgs.append(dcache.DcacheRequest(addVisitor=dcache.CaptchaID(id=captcha_id)))
#pipeline_msgs = dcache.DcacheBatchRequest(requests=pipeline_msgs)

#def pipeline_generate_messages():
#    for msg in pipeline_msgs:
#        yield msg


#import threading
#threadLock = threading.Lock()
#global_counter = 0

def nothing(r):
    pass


class HelloGrpcUser(adaptor.GrpcUser):
    stub_class = DcacheServiceStub
    host = host
    captcha_id = captcha_id
    msg = dcache.CaptchaID(id=captcha_id)
#    pipeline_msgs = []
#    for _ in range(0,10):
#        pipeline_msgs.append(dcache.DcacheRequest(addVisitor=dcache.CaptchaID(id=captcha_id)))
#    pipeline_msgs = dcache.DcacheBatchRequest(requests=pipeline_msgs)


#    pipeline_vote = []
#    for _ in range(0,1000):
#        pipeline_vote.append(
#        dcache.DcacheRequest(addVisitor=dcache.CaptchaID(id=captcha_id))
#        )




    def add_vote(self, captcha_id: str):
        resp = self.stub.AddVisitor(self.msg)

#    def add_vote_pipeline(self):
#        with threadLock:
#            global global_counter
#            global_counter += 1
#            print(f"count: {global_counter}")
        
#        res = self.stub.PipelineDcacheOps(pipeline_msgs)
#        for r in res.responses:
#            print(r.option_add_visitor_result.result.duration)
#       resp = self.pipeline_write(data=self.pipeline_vote)
#       #pprint(resp)

    @task
    def addVote(self):
        self.add_vote(self.captcha_id)

#    @task
#    def addVotePipeline(self):
#        self.add_vote_pipeline()
        #self.add_vote(self.captcha_id)
        #self.stub.PipelineDcacheOps(pipeline_generate_messages()):


#        msg = dcache.CaptchaID(id=captcha_id)
#        resp = self.stub.AddVisitor(msg)
#        pprint(resp)