📚 编程进阶教程站

第 13 章 · 测试非 HTTP 协议 ​

本章目标：突破 HTTP 限制，用 Locust 的自定义 User 类测试 TCP Socket、gRPC、XML-RPC 等协议，并理解 gevent monkey-patching 的约束条件。

13.1 扩展原理：request 事件 ​

Locust 原生只支持 HTTP/HTTPS，但架构上留了扩展口子。核心思路：

  1. 继承 User 类（而非 HttpUser）创建自定义客户端；
  2. 在每次协议调用完成后，手动触发 request 事件让 Locust 记录统计。
python
from locust import User, task, between
from locust.exception import LocustError
import socket
import time

class TcpClient:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.socket = None

    def connect(self):
        self.socket = socket.create_connection((self.host, self.port))

    def send_and_receive(self, payload: bytes) -> bytes:
        start_time = time.time()
        start_perf = time.perf_counter()
        try:
            self.socket.sendall(payload)
            response = self.socket.recv(4096)
            elapsed = time.perf_counter() - start_perf
            # 触发 request 事件让 Locust 记录
            events.request.fire(
                request_type="TCP",
                name="send_recv",
                response_time=elapsed * 1000,
                response_length=len(response),
                exception=None,
                context={},
            )
            return response
        except Exception as e:
            elapsed = time.perf_counter() - start_perf
            events.request.fire(
                request_type="TCP", name="send_recv",
                response_time=elapsed * 1000,
                response_length=0, exception=e, context={},
            )
            raise

    def close(self):
        if self.socket:
            self.socket.close()

13.2 gevent Monkey-Patching 约束 ​

关键限制

Locust 底层使用 gevent 协作式并发。所有被压测的协议库必须能被 gevent monkey-patch——纯 Python（用标准 socket 模块）的库开箱即用；但 C 扩展库（如 psycopg2）无法被 patch，会阻塞整个进程。

库可否 patch替代方案
socket / ssl✅ 开箱即用—
requests / urllib3✅ 已内置—
pymysql / asyncpg✅ 纯 Python—
psycopg2 (PostgreSQL)❌ C 扩展用 psycogreen 补丁
redis-py (hiredis)⚠️ 部分关闭 hiredis 解析器

13.3 XML-RPC 压测示例 ​

官方文档提供了 XML-RPC 客户端封装示例。先启动一个测试服务器：

python
# server.py
from xmlrpc.server import SimpleXMLRPCServer
import time, random

def get_time():
    time.sleep(random.random())
    return time.time()

def get_random_number(low, high):
    time.sleep(random.random())
    return random.randint(low, high)

server = SimpleXMLRPCServer(("localhost", 8877))
server.register_function(get_time, "get_time")
server.register_function(get_random_number, "get_random_number")
print("Listening on port 8877...")
server.serve_forever()

编写对应的 Locust User：

python
from xmlrpc.client import ServerProxy, Fault
import time

class XmlRpcClient(ServerProxy):
    def __init__(self, host, request_events):
        super().__init__(host)
        self._events = request_events

    def __getattr__(self, name):
        func = ServerProxy.__getattr__(self, name)
        def wrapper(*args, **kwargs):
            start = time.perf_counter()
            try:
                result = func(*args, **kwargs)
                elapsed_ms = (time.perf_counter() - start) * 1000
                self._events.request.fire(
                    request_type="XML-RPC", name=name,
                    response_time=elapsed_ms,
                    response_length=0,
                    exception=None, context={},
                )
                return result
            except Fault as e:
                elapsed_ms = (time.perf_counter() - start) * 1000
                self._events.request.fire(
                    request_type="XML-RPC", name=name,
                    response_time=elapsed_ms,
                    response_length=0,
                    exception=e, context={},
                )
                raise
        return wrapper

class XmlRpcUser(User):
    wait_time = between(1, 3)

    def on_start(self):
        self.client = XmlRpcClient("http://localhost:8877", environment.events)

    @task
    def call_get_time(self):
        self.client.get_time()

    @task(weight=3)
    def call_random(self):
        self.client.get_random_number(1, 100)

运行 python server.py & 然后 locust -f locustfile.py --headless -u 10 -r 2。

13.4 gRPC 测试思路 ​

gRPC 基于 HTTP/2 + Protobuf，需要 grpcio 和 grpcio-tools：

python
import grpc
import hello_pb2, hello_pb2_grpc  # 由 protoc 编译生成

class GrpcUser(User):
    wait_time = between(0.5, 2)

    def on_start(self):
        self.channel = grpc.insecure_channel("localhost:50051")
        self.stub = hello_pb2_grpc.GreeterStub(self.channel)

    @task
    def say_hello(self):
        start = time.perf_counter()
        try:
            response = self.stub.SayHello(
                hello_pb2.HelloRequest(name="locust"),
                timeout=5.0
            )
            elapsed = (time.perf_counter() - start) * 1000
            environment.events.request.fire(
                request_type="gRPC", name="SayHello",
                response_time=elapsed,
                response_length=len(response.message),
                exception=None, context={},
            )
        except grpc.RpcError as e:
            elapsed = (time.perf_counter() - start) * 1000
            environment.events.request.fire(
                request_type="gRPC", name="SayHello",
                response_time=elapsed,
                response_length=0,
                exception=e, context={},
            )

gevent 与 grpcio

grpcio 使用 C 核心，gevent 无法直接 patch。解决方案是给每个 User 分配独立 channel 并控制并发数，或考虑使用第 14 章介绍的 aiolocust。

本章小结 ​

  • 自定义 User 类继承 User（非 HttpUser），在调用后触发 environment.events.request.fire()；
  • 所有协议库必须可被 gevent monkey-patch——纯 Python 可以，C 扩展不行；
  • XML-RPC、Socket、MQTT 等均可按此模式封装；
  • gRPC 因 C 核心受限，建议评估 aiolocust 或控制并发。

🧪 随堂测验

点击你认为正确的选项。答错时会展示正确答案与原因解析。

1. 让 Locust 记录非 HTTP 协议调用的机制是什么？

2. 为什么 psycopg2 不能直接在 Locust 中使用？

3. 自定义非 HTTP User 类应该继承哪个基类？

4. 触发 Locust 统计事件的正确方法是？

🛠️ 动手实践 ​

  1. 编写一个简单的 TCP echo server，然后用自定义 User 类对其进行压测。
  2. 尝试用 pymysql 对 MySQL 数据库执行 SELECT 查询压测。
  3. 研究 psycogreen 文档，尝试让 psycopg2 与 gevent 兼容。

下一章探索 asyncio 异步方案——绕过 gevent 限制的现代路线。