From 8a2e118ab84168ea49206e1e548e9b6b61008dc1 Mon Sep 17 00:00:00 2001 From: xfy Date: Thu, 16 Apr 2026 11:44:45 +0800 Subject: [PATCH] =?UTF-8?q?test(proxy):=20=E6=B7=BB=E5=8A=A0=20WebSocket?= =?UTF-8?q?=20=E4=BB=A3=E7=90=86=E6=80=A7=E8=83=BD=E5=9F=BA=E5=87=86?= =?UTF-8?q?=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 覆盖握手升级请求构建、不同帧大小转发吞吐量和心跳往返延迟。 Co-Authored-By: Claude Opus 4.6 --- internal/proxy/websocket_bench_test.go | 397 +++++++++++++++++++++++++ 1 file changed, 397 insertions(+) create mode 100644 internal/proxy/websocket_bench_test.go diff --git a/internal/proxy/websocket_bench_test.go b/internal/proxy/websocket_bench_test.go new file mode 100644 index 0000000..d9300fb --- /dev/null +++ b/internal/proxy/websocket_bench_test.go @@ -0,0 +1,397 @@ +// Package proxy 提供 WebSocket 代理模块的性能基准测试。 +// +// 该文件测试 WebSocket 代理模块各项操作的性能，包括： +// - 握手升级请求构建性能 +// - 不同帧大小的数据转发性能 +// - Ping/Pong 心跳往返延迟 +// - 并发连接吞吐量 +// +// 注意：WebSocket 代理在 TCP 层进行透明转发，不解析 WebSocket 帧协议。 +// 帧编码/解码开销由底层 TCP 数据传输的编解码开销代表。 +// +// 作者：xfy +package proxy + +import ( + "net" + "net/http" + "strings" + "sync" + "testing" + + "github.com/valyala/fasthttp" +) + +// BenchmarkWebSocketHandshake 基准测试 WebSocket 握手升级请求构建性能。 +// +// 测试 buildWebSocketUpgradeRequest() 函数的开销，包括： +// - 路径和查询参数处理 +// - 请求头复制（Upgrade、Connection、Sec-WebSocket-*） +// - X-Forwarded 代理头注入 +// +// 模拟真实客户端请求，包含常见的 WebSocket 握手头。 +func BenchmarkWebSocketHandshake(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + ctx := &fasthttp.RequestCtx{} + ctx.Request.SetRequestURI("/ws?token=abc123&channel=default") + ctx.Request.Header.SetHost("client.example.com") + ctx.Request.Header.Set("Upgrade", "websocket") + ctx.Request.Header.Set("Connection", "Upgrade") + ctx.Request.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") + ctx.Request.Header.Set("Sec-WebSocket-Version", "13") + ctx.Request.Header.Set("Sec-WebSocket-Protocol", "chat, superchat") + ctx.Request.Header.Set("Sec-WebSocket-Extensions", "permessage-deflate") + ctx.Request.Header.Set("Origin", "https://example.com") + + result := buildWebSocketUpgradeRequest(ctx, "backend.example.com:8080") + + // 验证握手请求包含关键头 + if !strings.Contains(result, "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==") { + b.Fatal("handshake missing Sec-WebSocket-Key") + } + if !strings.Contains(result, "Upgrade: websocket") { + b.Fatal("handshake missing Upgrade header") + } + } +} + +// BenchmarkWebSocketFrameSmall 基准测试小帧（<126 bytes）发送接收性能。 +// +// 测试 copyData() 函数处理小数据块的开销。小帧对应 WebSocket +// 协议中 length < 126 的场景，帧头仅需 2 bytes（无 extended payload）。 +// 这是最常见的 WebSocket 帧大小（聊天消息、心跳、状态更新等）。 +func BenchmarkWebSocketFrameSmall(b *testing.B) { + smallData := make([]byte, 100) // 100 bytes < 126 + for i := range smallData { + smallData[i] = byte(i % 256) + } + + b.ReportAllocs() + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + src1, src2 := net.Pipe() + dst1, dst2 := net.Pipe() + + bridge := &WebSocketBridge{} + + // 启动数据复制 + done := make(chan struct{}) + go func() { + _ = bridge.copyData(dst1, src1, "bench-small") + close(done) + }() + + // 发送数据 + go func() { + _, _ = src2.Write(smallData) + _ = src2.Close() + }() + + // 读取数据 + buf := make([]byte, len(smallData)) + n, err := dst2.Read(buf) + if err != nil { + b.Fatalf("read error: %v", err) + } + if n != len(smallData) { + b.Fatalf("expected %d bytes, got %d", len(smallData), n) + } + + // 清理 + _ = dst1.Close() + _ = dst2.Close() + <-done + } + }) +} + +// BenchmarkWebSocketFrameMedium 基准测试中等帧（126-65535 bytes）发送接收性能。 +// +// 测试 copyData() 函数处理中等数据块的开销。中等帧对应 WebSocket +// 协议中 126 <= length < 65536 的场景，帧头需要 4 bytes +// （2 bytes 基础头 + 2 bytes extended payload length）。 +// 典型场景：JSON payloads、小型文件传输、表单数据。 +func BenchmarkWebSocketFrameMedium(b *testing.B) { + mediumData := make([]byte, 32*1024) // 32KB，典型中等帧大小 + for i := range mediumData { + mediumData[i] = byte(i % 256) + } + + b.ReportAllocs() + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + src1, src2 := net.Pipe() + dst1, dst2 := net.Pipe() + + bridge := &WebSocketBridge{} + + done := make(chan struct{}) + go func() { + _ = bridge.copyData(dst1, src1, "bench-medium") + close(done) + }() + + go func() { + _, _ = src2.Write(mediumData) + _ = src2.Close() + }() + + // 循环读取完整数据（可能分多次 Read） + received := 0 + buf := make([]byte, len(mediumData)) + for received < len(mediumData) { + n, err := dst2.Read(buf[received:]) + if err != nil { + b.Fatalf("read error at offset %d: %v", received, err) + } + received += n + } + if received != len(mediumData) { + b.Fatalf("expected %d bytes, got %d", len(mediumData), received) + } + + _ = dst1.Close() + _ = dst2.Close() + <-done + } + }) +} + +// BenchmarkWebSocketFrameLarge 基准测试大帧（>65535 bytes）发送接收性能。 +// +// 测试 copyData() 函数处理大数据块的开销。大帧对应 WebSocket +// 协议中 length >= 65536 的场景，帧头需要 10 bytes +// （2 bytes 基础头 + 8 bytes extended payload length）。 +// 典型场景：文件传输、视频帧、大数据集推送。 +// +// 使用 32KB 内部缓冲区，大帧需要多次 Read/Write 循环。 +func BenchmarkWebSocketFrameLarge(b *testing.B) { + largeData := make([]byte, 128*1024) // 128KB > 65535 + for i := range largeData { + largeData[i] = byte(i % 256) + } + + b.ReportAllocs() + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + src1, src2 := net.Pipe() + dst1, dst2 := net.Pipe() + + bridge := &WebSocketBridge{} + + done := make(chan struct{}) + go func() { + _ = bridge.copyData(dst1, src1, "bench-large") + close(done) + }() + + go func() { + _, _ = src2.Write(largeData) + _ = src2.Close() + }() + + // 循环读取完整数据 + received := 0 + buf := make([]byte, len(largeData)) + for received < len(largeData) { + n, err := dst2.Read(buf[received:]) + if err != nil { + b.Fatalf("read error at offset %d: %v", received, err) + } + received += n + } + if received != len(largeData) { + b.Fatalf("expected %d bytes, got %d", len(largeData), received) + } + + _ = dst1.Close() + _ = dst2.Close() + <-done + } + }) +} + +// BenchmarkWebSocketPingPong 基准测试 Ping/Pong 心跳往返性能。 +// +// WebSocket 代理在 TCP 层透明转发，Ping/Pong 控制帧由底层 +// TCP 连接直接传递。本测试模拟 Ping/Pong 的往返延迟： +// - 发送小数据包（Ping 帧，通常 0-125 bytes payload） +// - 等待对端接收并返回（Pong 帧） +// - 测量完整往返时间 +// +// 使用 16-byte 模拟 Ping payload，这是最常见的心跳大小。 +func BenchmarkWebSocketPingPong(b *testing.B) { + pingPayload := []byte("ping-1234567890") // 16 bytes，模拟 Ping 帧 + + b.ReportAllocs() + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + client1, client2 := net.Pipe() + target1, target2 := net.Pipe() + + // 启动双向桥接 + bridge := NewWebSocketBridge(client1, target1) + bridgeDone := make(chan struct{}) + go func() { + _ = bridge.Bridge() + close(bridgeDone) + }() + + // 客户端发送 Ping + go func() { + _, _ = client2.Write(pingPayload) + }() + + // 后端读取 Ping 并回传 Pong + buf := make([]byte, len(pingPayload)) + n, err := target2.Read(buf) + if err != nil { + b.Fatalf("target read error: %v", err) + } + if n != len(pingPayload) { + b.Fatalf("expected %d bytes, got %d", len(pingPayload), n) + } + + // 后端回传 Pong + go func() { + _, _ = target2.Write(buf[:n]) + }() + + // 客户端读取 Pong + pongBuf := make([]byte, len(pingPayload)) + n, err = client2.Read(pongBuf) + if err != nil { + b.Fatalf("client read error: %v", err) + } + if n != len(pingPayload) { + b.Fatalf("expected %d bytes Pong, got %d", len(pingPayload), n) + } + + // 清理 + _ = client2.Close() + _ = target2.Close() + <-bridgeDone + } + }) +} + +// BenchmarkWebSocketConcurrent 基准测试并发连接吞吐量。 +// +// 模拟多个 WebSocket 连接同时进行双向数据转发的场景。 +// 测试 Bridge() 函数在高并发下的性能表现，包括： +// - 多个 goroutine 同时读写 +// - sync.Mutex 保护关闭状态的开销 +// - 双向数据转发的总吞吐量 +// +// 每个连接发送 4KB 数据，模拟典型的 WebSocket 消息大小。 +func BenchmarkWebSocketConcurrent(b *testing.B) { + messageSize := 4 * 1024 // 4KB 典型消息大小 + message := make([]byte, messageSize) + for i := range message { + message[i] = byte(i % 256) + } + + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + client1, client2 := net.Pipe() + target1, target2 := net.Pipe() + + bridge := NewWebSocketBridge(client1, target1) + + // 启动桥接 + bridgeDone := make(chan struct{}) + go func() { + _ = bridge.Bridge() + close(bridgeDone) + }() + + var wg sync.WaitGroup + wg.Add(2) + + // 客户端 -> 后端 + go func() { + defer wg.Done() + _, _ = client2.Write(message) + }() + + // 后端读取客户端数据 + go func() { + defer wg.Done() + buf := make([]byte, messageSize) + received := 0 + for received < messageSize { + n, err := target2.Read(buf[received:]) + if err != nil { + break + } + received += n + } + }() + + wg.Wait() + + // 清理 + _ = client2.Close() + _ = target2.Close() + <-bridgeDone + } +} + +// BenchmarkWebSocketWriteUpgradeResponse 基准测试升级响应写入性能。 +// +// 测试 writeUpgradeResponse() 函数将 HTTP 101 响应写回客户端的开销。 +// 包含响应行构建、头格式化和单次 Write 调用。 +func BenchmarkWebSocketWriteUpgradeResponse(b *testing.B) { + resp := &http.Response{ + ProtoMajor: 1, + ProtoMinor: 1, + Status: "101 Switching Protocols", + StatusCode: 101, + Header: http.Header{ + "Upgrade": []string{"websocket"}, + "Connection": []string{"Upgrade"}, + "Sec-WebSocket-Accept": []string{"s3pPLMBiTxaQ9kYGzzhZRbK+xOo="}, + }, + } + + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + conn1, conn2 := net.Pipe() + + done := make(chan error, 1) + go func() { + done <- writeUpgradeResponse(conn1, resp) + _ = conn1.Close() + }() + + // 读取响应 + buf := make([]byte, 1024) + n, err := conn2.Read(buf) + if err != nil { + b.Fatalf("read error: %v", err) + } + + response := string(buf[:n]) + if !strings.Contains(response, "101 Switching Protocols") { + b.Fatal("response missing status line") + } + + _ = conn2.Close() + <-done + } +}