locust模擬tcp客戶端進行socket並發

 

一、背景

　　打算使用locust做並發，實現性能測試，網上找到的大部分都是http協議的測試，但被測平台是類似於IOT的平台，壓力不在於http的訪問，而在於終端設備的連接與數據發送，於是就想基於locust中http的使用方式，做一個tcp的並發腳本。

二、HttpUser

　　locust執行並發的腳本大致可分為兩部分，TaskSet和User。而HttpUser是繼承於User類，實現了http訪問相關的一些方法，只要繼承該類就可以很方便地進行並發測試，此時用戶只需要關注業務的實現即可，業務中方法的調用也和requests模塊的一致。

　　我們發現繼承自TaskSet的類中，可以直接調用 self.client 方法進行http的訪問，於是追蹤 self.client 方法的源碼。　　

    @property
    def client(self):
        """
        Shortcut to the client :py:attr:`client <locust.User.client>` attribute of this TaskSet's :py:class:`User <locust.User>`
        """
        return self.user.client

　　返回的是 self.user.client，繼續追蹤 self.user

    @property
    def user(self):
        """:py:class:`User <locust.User>` instance that this TaskSet was created by"""
        return self._user

　　繼續追蹤 self._user

    def __init__(self, parent):
        self._task_queue = []
        self._time_start = time()

        if isinstance(parent, TaskSet):
            self._user = parent.user
        else:
            self._user = parent

        self._parent = parent

　　在繼承自TaskSet的類中打印 self.parent 可以發現，他就是繼承自 HttpUser的類，這時候關鍵點就來到了HttpUser，此時可以認為self.client就是HttpUser.client

　　點進HttpUser查看源碼。

class HttpUser(User):
    """
    Represents an HTTP "user" which is to be spawned and attack the system that is to be load tested.

    The behaviour of this user is defined by its tasks. Tasks can be declared either directly on the
    class by using the :py:func:`@task decorator <locust.task>` on methods, or by setting
    the :py:attr:`tasks attribute <locust.User.tasks>`.

    This class creates a *client* attribute on instantiation which is an HTTP client with support
    for keeping a user session between requests.
    """

    abstract = True
    """If abstract is True, the class is meant to be subclassed, and users will not choose this locust during a test"""

    client: HttpSession = None
    """
    Instance of HttpSession that is created upon instantiation of Locust.
    The client supports cookies, and therefore keeps the session between HTTP requests.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if self.host is None:
            raise LocustError(
                "You must specify the base host. Either in the host attribute in the User class, or on the command line using the --host option."
            )

        session = HttpSession(
            base_url=self.host,
            request_event=self.environment.events.request,
            user=self,
        )
        session.trust_env = False
        self.client = session

　　發現 HttpUser確實有個client屬性，該屬性是一個HttpSession類的實例化對象。

　　點進HttpSession類，可以發現，該類是繼承自requests.Session，而request正是實現http請求的常用模塊。

三、TcpSocketUser構造

　　至此，我們可以大膽猜測，是不是只要構造一個 TcpSocketUser繼承自User，並且將一個實現了socket方法的實例化對象賦值給TcpSocketUser.client方法即可呢？

# locustfile.py
# coding:utf-8
import socket
from locust import User, TaskSet, task
import time


class TcpSocketUser(User):

    abstract = True

    client: socket.socket = None

    def __init__(self, *args, **kwargs):
        super(TcpSocketUser, self).__init__(*args, **kwargs)
        self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        host, port = self.host.split(":")
        self.client.connect((host, int(port)))


class TestTaskSet(TaskSet):

    @task(1)
    def send_data(self):
        while True:
            self.client.send(b"This is send data!")
            data = self.client.recv(1024)
            print(data)
            time.sleep(3)


class TcpTestUser(TcpSocketUser):

    tasks = [TestTaskSet]
    host = "172.16.100.139:8001"


if __name__ == "__main__":
    import os
    os.system("locust")

# Server.py
# coding:utf-8
import gevent
from gevent import socket, monkey

monkey.patch_all()


def server(port):
    try:
        s = socket.socket()
        s.bind(('0.0.0.0', port))
        s.listen(500)
        while True:
            cli, addr = s.accept()
            gevent.spawn(handle_request, cli)
    except KeyboardInterrupt as e:
        print(e)


def handle_request(conn):
    try:
        while True:
            data = conn.recv(1024)
            if not data:
                conn.close()
            else:
                print("recv:", data)
                conn.send(b"Received!")
    except OSError as e:
        print("client has been closed")

    except Exception as ex:
        print(ex)
    finally:
        conn.close()


if __name__ == '__main__':
    server(8001)

 　　以上的代碼實現了功能，客戶端定時向服務端發送數據，服務端收到后進行回復，運行后發現可行，此時初步的並發框架已經完成。

　　此時想實現第二個需求，客戶端每隔5s同時向服務器發送兩條數據 “This is First data!” 和 “This is Second data!”，但是服務端因為要處理的原因，只能20s后才能返回結果，實現代碼如下。

# locustfile.py
# coding:utf-8
import socket
from locust import User, TaskSet, task, constant
import time


class TcpSocketUser(User):

    abstract = True

    client: socket.socket = None

    def __init__(self, *args, **kwargs):
        super(TcpSocketUser, self).__init__(*args, **kwargs)
        self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        host, port = self.host.split(":")
        self.client.connect((host, int(port)))


class TestTaskSet(TaskSet):

    @task(1)
    def send_data(self):
        send_data = b"This is First data!"
        self.client.send(send_data)
        print(send_data)
        send_data = b"This is Second data!"
        self.client.send(send_data)
        print(send_data)
        recv_data = self.client.recv(1024)
        print(recv_data)


class TcpTestUser(TcpSocketUser):

    wait_time = constant(5)

    tasks = [TestTaskSet]
    host = "172.16.100.139:8001"


if __name__ == "__main__":
    import os
    os.system("locust")

# Server.py
# coding:utf-8
import gevent
from gevent import socket, monkey
import time

monkey.patch_all()


def server(port):
    try:
        s = socket.socket()
        s.bind(('0.0.0.0', port))
        s.listen(500)
        while True:
            cli, addr = s.accept()
            gevent.spawn(handle_request, cli)
    except KeyboardInterrupt as e:
        print(e)


def handle_request(conn):
    try:
        while True:
            data = conn.recv(1024)
            if not data:
                conn.close()
            else:
                print("recv:", data)
                time.sleep(20)
                if b"First" in data:
                    conn.send(b"First Received!")
                elif b'Second' in data:
                    conn.send(b"Second Received!")
    except OSError as e:
        print("client has been closed")

    except Exception as ex:
        print(ex)
    finally:
        conn.close()


if __name__ == '__main__':
    server(8001)

　　運行起來會發現，當客戶端數據發送完成之后，只能等待20s，收到服務端的答復后才能返回開始下一次數據發送，這是因為recv默認是個阻塞方法，當沒有收到數據時會一直阻塞在此，影響下一次任務的開始，此時我們要解決阻塞的問題，可以使用recv的block=False，但官方文檔在Testing non-HTTP systems 章節提到，協議庫要用可以猴子補丁的。

　　所以我們選擇使用gevent的方法來進行處理。

四、gevent介紹　

　　gevent是基於協程的Python網絡庫。特點：

• 基於libev的快速事件循環(Linux上epoll，FreeBSD上kqueue）。
• 基於greenlet的輕量級執行單元。
• API的概念和Python標准庫一致(如事件，隊列)。
• 可以配合socket，ssl模塊使用。
• 能夠使用標准庫和第三方模塊創建標准的阻塞套接字(gevent.monkey)。
• 默認通過線程池進行DNS查詢,也可通過c-are(通過GEVENT_RESOLVER=ares環境變量開啟）。
• TCP/UDP/HTTP服務器
• 子進程支持（通過gevent.subprocess）
• 線程池

　　我們主要是用到上面標注紅色的特點，該特點是結合猴子補丁一起使用的，猴子補丁的作用是將一些阻塞的方法/模塊動態替換成非阻塞的，比如 socket、time等。

五、非阻塞代碼　

# locustfile.py
# coding:utf-8
from gevent import socket, monkey;monkey.patch_all()
import gevent
import socket
from locust import User, TaskSet, task, constant
import time


class TcpSocketUser(User):

    abstract = True

    client: socket.socket = None

    def __init__(self, *args, **kwargs):
        super(TcpSocketUser, self).__init__(*args, **kwargs)
        self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        host, port = self.host.split(":")
        self.client.connect((host, int(port)))
        gevent.spawn(handle_recv, self.client)


class TestTaskSet(TaskSet):

    @task(1)
    def send_data(self):
        send_data = b"This is First data!"
        self.client.send(send_data)
        print(send_data)
        send_data = b"This is Second data!"
        self.client.send(send_data)
        print(send_data)


def handle_recv(conn):
    while True:
        data = conn.recv(1024)
        if not data:
            conn.close()
        else:
            print(data)
        time.sleep(0.1)


class TcpTestUser(TcpSocketUser):

    wait_time = constant(5)

    tasks = [TestTaskSet]
    host = "172.16.100.139:8001"


if __name__ == "__main__":
    import os
    os.system("locust")

# Server.py
# coding:utf-8
import gevent
from gevent import socket, monkey
import time

monkey.patch_all()


def server(port):
    try:
        s = socket.socket()
        s.bind(('0.0.0.0', port))
        s.listen(500)
        while True:
            cli, addr = s.accept()
            gevent.spawn(handle_request, cli)
    except KeyboardInterrupt as e:
        print(e)


def handle_request(conn):
    try:
        while True:
            data = conn.recv(1024)
            if not data:
                conn.close()
            else:
                print("recv:", data)
                time.sleep(20)
                for i in data.split(b'!'):
                    if b"First" in i:
                        conn.send(b"First Received!")
                    elif b'Second' in i:
                        conn.send(b"Second Received!")
    except OSError as e:
        print("client has been closed")

    except Exception as ex:
        print(ex)
    finally:
        conn.close()


if __name__ == '__main__':
    server(8001)

　　以上代碼可實現上述第二個需求，即 客戶端每隔5s同時向服務器發送兩條數據 “This is First data!” 和 “This is Second data!”，但是服務端因為要處理的原因，只能20s后才能返回結果。

　　至此，已經可以實現並發了。

六、並發結果注冊到測試報告

　　我們繼續回到HttpUser源碼中查看。

　　HttpSession實例化時傳入三個參數，self.host，self.environment.events.request，self，跟報告相關的主要是第二個參數，我們到HttpSession中進行查看。

　　發現兩個跟報告相關的地方。　　

        request_meta = {
            "request_type": method,
            "response_time": response_time,
            "name": name or url_after_redirect,
            "context": context,
            "response": response,
            "exception": None,
            "start_time": start_time,
            "url": request_after_redirect.url,
        }

        if catch_response:
            return ResponseContextManager(response, request_event=self.request_event, request_meta=request_meta)

　　request_meta中的各個key與報告中的每一列是一一對應的。

　　因此我們認為 ResponseContextManager 類中就含有將結果注冊到測試報告的相關代碼，於是我們繼續到 ResponseContextManager 源碼中查看，果然發現了相關的代碼。　

    def _report_request(self, exc=None):
        self._request_event.fire(**self.request_meta)

　　self._request_event指向的就是self.environment.events.request，所以fire方法就可以將request_meta的數據注冊到報告中。

　　此時我們做一個小調整，新建一個TcpSocketClient類繼承自socket類，重寫connect、recv、send等方法，使其可以注冊到報告中。

　　request_meta中大部分參數的含義都比較清晰，這里不做一一解釋（本文未用到response和context），但問題在於在注冊到報告中，如何確認是成功或失敗呢？

　　查看locust官方文檔，可以看到有這么一段。

Parameters
request_type – Request type method used

name – Path to the URL that was called (or override name if it was used in the call to the client)

response_time – Time in milliseconds until exception was thrown

response_length – Content-length of the response

response – Response object (e.g. a requests.Response)

context – User/request context

exception – Exception instance that was thrown. None if request was successful.

　　其中exception用法寫到，該參數傳入的是方法執行錯誤產生時所拋出的錯誤實例，如果方法執行成功則填為None，至此，大部分代碼都已完成，我們看下代碼和結果。

# locustfile.py
# coding:utf-8
from gevent import socket, monkey, event;monkey.patch_all()
import gevent
import socket
from locust import User, TaskSet, task, constant
import time


class TcpSocketClient(socket.socket):

    def __init__(self, af_inet, socket_type, request_event, user):
        super(TcpSocketClient, self).__init__(af_inet, socket_type)
        self.connect_event = event.Event()
        self.request_event = request_event
        self.user = user

    def connect(self, addr):
        start_time = time.time()
        try:
            super(TcpSocketClient, self).connect(addr)
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            self.request_event.fire(request_type="tcp", name="connect", response_time=total_time, response_length=0,
                                    response=None, context=None, exception=e)
        else:
            self.connect_event.set()
            total_time = int((time.time() - start_time) * 1000)
            self.request_event.fire(request_type="tcp", name="connect", response_time=total_time, response_length=0,
                                    response=None, context=None, exception=None)

    def send(self, data, flag=False):
        start_time = time.time()
        try:
            self.connect_event.wait()
            super(TcpSocketClient, self).send(data)
            total_time = int((time.time() - start_time) * 1000)
            print(data)
            self.request_event.fire(request_type="tcp", name="send", response_time=total_time,
                                    response_length=len(data), response=None, context=None, exception=None)
        except OSError as e:
            total_time = int((time.time() - start_time) * 1000)
            self.request_event.fire(request_type="tcp", name="send", response_time=total_time,
                                    response_length=0, response=None, context=None, exception=e)

    def recv(self, bufsize, flag=False):
        start_time = time.time()
        try:
            while True:
                start_time = time.time()
                self.connect_event.wait()
                data = super(TcpSocketClient, self).recv(bufsize)
                total_time = int((time.time() - start_time) * 1000)
                print(data)
                self.request_event.fire(request_type="tcp", name="recv", response_time=total_time,
                                        response_length=len(data), response=None, context=None, exception=None)
                time.sleep(0.1)
        except OSError as e:
            total_time = int((time.time() - start_time) * 1000)
            self.request_event.fire(request_type="tcp", name="recv", response_time=total_time, response_length=0,
                                    response=None, context=None, exception=e)


class TcpSocketUser(User):

    abstract = True

    client: socket.socket = None

    def __init__(self, *args, **kwargs):
        super(TcpSocketUser, self).__init__(*args, **kwargs)
        self.client = TcpSocketClient(socket.AF_INET, socket.SOCK_STREAM, self.environment.events.request, self)
        host, port = self.host.split(":")
        self.client.connect((host, int(port)))
        gevent.spawn(self.client.recv, 1024)


class TestTaskSet(TaskSet):

    @task(1)
    def send_data(self):
        send_data = b"This is First data!"
        self.client.send(send_data)
        send_data = b"This is Second data!"
        self.client.send(send_data)


class TcpTestUser(TcpSocketUser):

    wait_time = constant(5)

    tasks = [TestTaskSet]
    host = "172.16.100.139:8001"


if __name__ == "__main__":
    import os
    os.system("locust")

# Server.py
# coding:utf-8
import gevent
from gevent import socket, monkey
import time

monkey.patch_all()


def server(port):
    try:
        s = socket.socket()
        s.bind(('0.0.0.0', port))
        s.listen(500)
        while True:
            cli, addr = s.accept()
            gevent.spawn(handle_request, cli)
    except KeyboardInterrupt as e:
        print(e)


def handle_request(conn):
    try:
        while True:
            data = conn.recv(1024)
            if not data:
                conn.close()
            else:
                print("recv:", data)
                # time.sleep(5)
                for i in data.split(b'!'):
                    if b"First" in i:
                        conn.send(b"First Received!")
                    elif b'Second' in i:
                        conn.send(b"Second Received!")
    except OSError as e:
        print("client has been closed")

    except Exception as ex:
        print(ex)
    finally:
        conn.close()


if __name__ == '__main__':
    server(8001)

七、業務融合　　

　　這里又出現一個問題，注冊到報告中的信息，更多是想關注業務的成功與失敗，而不是底層的connect、send、recv這些，那該怎么辦呢？

　　因為我們要模擬的是終端客戶向平台發送數據，所以我們新建一個Device類，該類繼承自TcpSocketClient，並且會擴展實際的業務。

　　現在假設我們有這么一個需求，客戶端連接平台后要先發送一條“Register！”指令向平台注冊，如果平台回復“Success”，則客戶端開始以5s的間隔持續向平台發送“Heart！”，若平台超過5s沒有響應，則客戶端注冊失敗，不再繼續發送信息。

　　此時我們更多的關注在於業務的響應時間，比如register和heart的回復時間。

　　實現代碼如下：　

# locustfile.py
# coding:utf-8
from gevent import socket, monkey, event;monkey.patch_all()
import gevent
import socket
from locust import User, TaskSet, task, constant
import time
import queue


class TcpSocketClient(socket.socket):

    def __init__(self, af_inet, socket_type, request_event, user):
        super(TcpSocketClient, self).__init__(af_inet, socket_type)
        self.connect_event = event.Event()
        self.request_event = request_event
        self.user = user
        self.recv_queue = queue.Queue()

    def connect(self, addr):
        start_time = time.time()
        try:
            super(TcpSocketClient, self).connect(addr)
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            self.request_event.fire(request_type="tcp", name="connect", response_time=total_time, response_length=0,
                                    response=None, context=None, exception=e)
        else:
            self.connect_event.set()
            total_time = int((time.time() - start_time) * 1000)
            self.request_event.fire(request_type="tcp", name="connect", response_time=total_time, response_length=0,
                                    response=None, context=None, exception=None)

    def send(self, data, flag=False):
        start_time = time.time()
        try:
            self.connect_event.wait()
            super(TcpSocketClient, self).send(data)
            total_time = int((time.time() - start_time) * 1000)
            print(self.getsockname(), data)
            self.request_event.fire(request_type="tcp", name="send", response_time=total_time,
                                    response_length=len(data), response=None, context=None, exception=None)
        except OSError as e:
            total_time = int((time.time() - start_time) * 1000)
            self.request_event.fire(request_type="tcp", name="send", response_time=total_time,
                                    response_length=0, response=None, context=None, exception=e)

    def recv(self, bufsize, flag=False):
        start_time = time.time()
        try:
            while True:
                start_time = time.time()
                self.connect_event.wait()
                data = super(TcpSocketClient, self).recv(bufsize)
                total_time = int((time.time() - start_time) * 1000)
                self.recv_queue.put_nowait(data)
                self.request_event.fire(request_type="tcp", name="recv", response_time=total_time,
                                        response_length=len(data), response=None, context=None, exception=None)
                time.sleep(0.1)
        except OSError as e:
            total_time = int((time.time() - start_time) * 1000)
            self.request_event.fire(request_type="tcp", name="recv", response_time=total_time, response_length=0,
                                    response=None, context=None, exception=e)


class Device(TcpSocketClient):
    def __init__(self, af_inet, socket_type, request_event, user):
        super(Device, self).__init__(af_inet, socket_type, request_event, user)
        self.is_registered = False
        self.heart_dict = {}

    def parse_data(self):
        while True:
            if not self.recv_queue.empty():
                data = self.recv_queue.get_nowait()
                print(self.getsockname(), data)
                if b"Success" in data:
                    self.is_registered = True
                elif b"Heart Reply" in data:
                    start_time = self.heart_dict.get("Heart")
                    total_time = int((time.time() - start_time) * 1000)
                    self.request_event.fire(request_type="tcp", name="Heart", response_time=total_time,
                                            response_length=0, response=None, context=None, exception=None)
            time.sleep(0.1)

    def register(self):
        self.send(b"Register!")

    def heart(self):
        self.send(b"Heart!")


class TcpSocketUser(User):

    abstract = True

    client: socket.socket = None

    def __init__(self, *args, **kwargs):
        super(TcpSocketUser, self).__init__(*args, **kwargs)
        self.client = Device(socket.AF_INET, socket.SOCK_STREAM, self.environment.events.request, self)
        host, port = self.host.split(":")
        self.client.connect((host, int(port)))
        gevent.spawn(self.client.recv, 1024)
        gevent.spawn(self.client.parse_data)


class TestTaskSet(TaskSet):

    def on_start(self):
        start_time = time.time()
        try:
            self.client.register()
            count = 0
            while not self.client.is_registered:
                time.sleep(0.1)
                count += 1
                if count == 100:
                    raise TimeoutError("注冊失敗")
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            self.client.request_event.fire(request_type="tcp", name="register", response_time=total_time,
                                           response_length=0, response=None, context=None, exception=e)
        else:
            total_time = int((time.time() - start_time) * 1000)
            self.client.request_event.fire(request_type="tcp", name="register", response_time=total_time,
                                           response_length=0, response=None, context=None, exception=None)

    @task(1)
    def heart(self):
        if self.client.is_registered:
            self.client.heart()
            self.client.heart_dict["Heart"] = time.time()


class TcpTestUser(TcpSocketUser):

    wait_time = constant(5)

    tasks = [TestTaskSet]
    host = "172.16.100.139:8001"


if __name__ == "__main__":
    import os
    os.system("locust")

# Server.py
# coding:utf-8
import gevent
from gevent import socket, monkey
import time

monkey.patch_all()


def server(port):
    try:
        s = socket.socket()
        s.bind(('0.0.0.0', port))
        s.listen(500)
        while True:
            cli, addr = s.accept()
            gevent.spawn(handle_request, cli)
    except KeyboardInterrupt as e:
        print(e)


def handle_request(conn):
    try:
        while True:
            data = conn.recv(1024)
            if not data:
                conn.close()
            else:
                print("recv:", data)
                # time.sleep(5)
                for i in data.split(b'!'):
                    if b"Register" in i:
                        conn.send(b"Success!")
                        pass
                    if b"Heart" in i:
                        conn.send(b"Heart Reply!")

    except OSError as e:
        print("client has been closed")

    except Exception as ex:
        print(ex)
    finally:
        conn.close()


if __name__ == '__main__':
    server(8001)

　　

 

　　至此，一個模擬客戶端並發連接平台的測試腳本案例完成。