求基于Locust的gRPC双向通信性能测试可运行示例

前置依赖 #

先安装必要的Python包：

pip install grpcio grpcio-tools locust

1. 定义gRPC双向流服务 #

创建chat.proto文件，定义双向流式通信的服务：

syntax = "proto3";

package chat;

service ChatService {
  rpc Chat(stream Message) returns (stream Message) {}
}

message Message {
  string user = 1;
  string content = 2;
}

执行以下命令生成Python代码：

python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. chat.proto

执行后会生成chat_pb2.py和chat_pb2_grpc.py两个文件。

2. 示例gRPC服务端（用于测试） #

创建grpc_server.py，实现一个简单的双向流echo服务：

import asyncio
import grpc
from chat_pb2 import Message
from chat_pb2_grpc import ChatServiceServicer, add_ChatServiceServicer_to_server

class ChatService(ChatServiceServicer):
    async def Chat(self, request_iterator, context):
        async for msg in request_iterator:
            yield Message(user=f"Server", content=f"Received: {msg.content}")

async def serve():
    server = grpc.aio.server()
    add_ChatServiceServicer_to_server(ChatService(), server)
    listen_addr = "[::]:50051"
    server.add_insecure_port(listen_addr)
    print(f"Starting server on {listen_addr}")
    await server.start()
    await server.wait_for_termination()

if __name__ == "__main__":
    asyncio.run(serve())

3. Locust性能测试脚本 #

创建locustfile.py，实现基于Locust的双向流性能测试：

import asyncio
from locust import User, task, events
import grpc
from chat_pb2 import Message
from chat_pb2_grpc import ChatServiceStub

class GrpcUser(User):
    abstract = True
    stub = None

    def on_start(self):
        # 建立异步gRPC通道
        self.channel = grpc.aio.insecure_channel("localhost:50051")
        self.stub = ChatServiceStub(self.channel)

    def on_stop(self):
        # 关闭通道
        asyncio.create_task(self.channel.close())

class ChatUser(GrpcUser):
    @task
    async def chat_stream(self):
        async def send_messages(stream):
            # 模拟发送3条消息
            for i in range(3):
                msg = Message(user=f"User_{self.id}", content=f"Message_{i}")
                await stream.write(msg)
                await asyncio.sleep(0.5)
            await stream.done_writing()

        try:
            # 调用双向流方法
            stream = self.stub.Chat()
            # 启动发送消息的任务
            send_task = asyncio.create_task(send_messages(stream))
            # 接收服务端响应
            async for response in stream:
                print(f"Received from server: {response.content}")
            await send_task
        except grpc.RpcError as e:
            events.request_failure.fire(
                request_type="grpc",
                name="ChatStream",
                response_time=0,
                exception=e
            )
        else:
            events.request_success.fire(
                request_type="grpc",
                name="ChatStream",
                response_time=0,
                response_length=0
            )

4. 运行步骤 #

1. 启动gRPC服务端：

   python grpc_server.py
   

2. 启动Locust：

   locust -f locustfile.py
   

3. 打开浏览器访问http://localhost:8089，设置并发用户数、孵化率，点击"Start swarming"开始测试。

关键说明 #

• 采用异步gRPC与Locust的异步任务模型适配，避免阻塞Locust的事件循环
• on_start/on_stop方法处理gRPC通道的创建与销毁，复用连接提升性能
• 任务中同时处理消息发送和接收逻辑，模拟真实双向流场景
• 通过events模块上报请求的成功/失败状态，便于Locust统计指标

内容的提问来源于stack exchange，提问作者user2416212