如何在 locust 中创建自定义套接字客户端

Question

I want to custom a tcp socket client in locust.And before add event ,my code was worked.

locustTcp.py

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

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)

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)
        ADDR = ('127.0.0.1', 9999)
        self.client.connect(ADDR)


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

    @task
    def send_data(self):
        msg = "hello"
        self.client.send(msg.encode())
        data = self.client.recv(1024).decode()
        print(data)

Server.py

#-*- 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()
        time.sleep(1)
        print('new connect from {}'.format(addr))
        t = threading.Thread(target=tcplink, args=(sock, addr))
        t.start()

main()

And then I add a warpper func in getattribute in TcpSocketClient class just like the docs page https://docs.locust.io/en/stable/testing-other-systems.html?

So my TcpSocketClient become to

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 __getattribute__(self, item):
        func = socket.socket.__getattribute__(self, item)

        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)
                self._locust_environment.events.request_success.fire(request_type="tcpsocket", name=item,
                                                                     response_time=total_time, response_length=0)
            return wrapper

And then I enter locust -f locustTcp.py will rise a error

Traceback (most recent call last):
  File "src\\gevent\\greenlet.py", line 854, in gevent._gevent_cgreenlet.Greenlet.run
  File "d:\test-project\venv\lib\site-packages\locust\runners.py", line 417, in <lambda>
    lambda: super(LocalRunner, self).start(user_count, spawn_rate, wait=wait)
  File "d:\test-project\venv\lib\site-packages\locust\runners.py", line 299, in start
    self.spawn_users(user_count, spawn_rate=spawn_rate, wait=wait)
  File "d:\test-project\venv\lib\site-packages\locust\runners.py", line 194, in spawn_users
    spawn()
  File "d:\test-project\venv\lib\site-packages\locust\runners.py", line 187, in spawn
    new_user = user_class(self.environment)
  File "D:\test-project\locustTcp.py", line 71, in __init__
    self.client = TcpSocketClient(socket.AF_INET, socket.SOCK_STREAM)
  File "D:\test-project\locustTcp.py", line 13, in __init__
    super().__init__(af_inet, socket_type)
  File "d:\test-project\venv\lib\site-packages\gevent\_socket3.py", line 143, in __init__
    self._sock = self._gevent_sock_class(family, type, proto, fileno)
TypeError: 'NoneType' object is not callable

So how can I add the event successfully to costom a socket client? Thanks a lot! My python verion is 3.7.5 in windows10

Answer 1

I'm having trouble understanding the callstack, but you dont seem to set _locust_environment in TcpSocketClient. In the example in the documentation, the last row in the User constructor is:

self.client._locust_environment = self.environment

Maybe that is what is missing? If not, try running it in a debugger to see where it goes wrong.

If you dont know how run in the debugger, check https://github.com/SvenskaSpel/locust-plugins/blob/master/examples/debug_ex.py