locust 新手問下 locust 自定義一個 socket 客戶端該如何實現

denki發表於2020-09-11

最近想學習一下 locust 跑一下壓測,locust 本身不支援 socket,所以打算按照文件裡自定義 xml-rpc 的方法自己先寫一個簡單的試試手。發現客戶端在 getattr 里加入 wrapper 後執行會丟擲一個 Nonetype 的錯誤,找了很久資料不知道如何解決,希望 testhome 的各位大佬能給點意見,如果有簡易版的實現方法能讓我參考一下就更好了。谷歌到的要麼是 websocket 的要麼是完全不能用的程式碼。
服務端程式碼

#-*- coding:utf-8 -*-

import time
import socket
import threading


def tcplink(sock, addr):
    msg = 'welcome'
    sock.send(msg.encode())
    while True:
        try:
            data = sock.recv(1024).decode()
            time.sleep(1)
            if data == 'exit' or not data:
                break
            sock.send('hello {}'.format(data).encode())
        except Exception as e:
            print(e)
            break
    sock.close()
    print('connect close from {}'.format(addr))


def main():
    addr = ('127.0.0.1', 9999)
    tctime = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tctime.bind(addr)
    tctime.listen(3)
    print('waiting for connection')
    while True:
        sock, addr = tctime.accept()
        print('new connect from {}'.format(addr))
        t = threading.Thread(target=tcplink, args=(sock, addr))
        t.start()
main()

客戶端程式碼

import time
import random
import socket
from locust import User, TaskSet, events, task, event, HttpUser
from locust_plugins import run_single_user


class TcpSocketClient(socket.socket):

    _locust_environment = None

    def __init__(self, af_inet, socket_type):
        print('start init tcp socekt')
        super().__init__(af_inet, socket_type)

    def __getattr__(self, item):
        func = socket.socket.__getattr__(self, item)
        print('is it tcp wrapper')

        def wrapper(*args, **kwargs):
            start_time = time.time()
            try:

                result = func(*args, **kwargs)
            except Exception as e:
                total_time = int((time.time() - start_time) * 1000)
                self._locust_environment.events.request_failure.fire(request_type="tcpsocket", name=item,
                                                                     response_time=total_time, exception=e)
            else:
                total_time = int((time.time() - start_time) * 1000)
                print(total_time, 'total time is ')
                self._locust_environment.events.request_success.fire(request_type="tcpsocket", name=item,
                                                                    response_time=total_time, response_length=0)
                return result
            return wrapper


class TcpSocketLocust(User):
    abstract = True

    def __init__(self, *args, **kwargs):
        print('it is tcpsocket locust')
        super().__init__(*args, **kwargs)
        self.client = TcpSocketClient(socket.AF_INET, socket.SOCK_STREAM)
        self.client._locust_environment = self.environment
        ADDR = ('127.0.0.1', 9999)
        self.client.connect(ADDR)


class TcpTestUser(TcpSocketLocust):
    min_wait = 10
    max_wait = 20

    @task
    def send_data(self):
        msg = "User"

        self.client.send(msg.encode())

        data = self.client.recv(1024).decode()
        print(data)


if __name__ == '__main__':
    run_single_user(TcpTestUser)

丟擲的 error

Traceback (most recent call last):
  File "D:/test-project/locustTcp.py", line 64, in <module>
    run_single_user(TcpTestUser)
  File "D:\test-project\venv\lib\site-packages\locust_plugins\debug.py", line 61, in run_single_user
    locust_class(env).run()
  File "D:/test-project/locustTcp.py", line 46, in __init__
    self.client.connect(ADDR)
  File "D:\test-project\venv\lib\site-packages\gevent\_socket3.py", line 405, in connect
    err = self.getsockopt(SOL_SOCKET, SO_ERROR)
TypeError: 'NoneType' object is not callable

相關文章