4. 编写Go压测脚本：使用vegeta库编写自定义压测程序、控制并发与速率

在生产级尾延迟诊断中，直接使用 vegeta 命令行工具虽然方便，但无法灵活控制请求构造、动态参数、自定义指标采集和退出条件。本章将深入 vegeta 的 Go 库 API，手写一个可定制、可观测的压测脚本，并精确控制并发数与请求速率。

4.1 为什么需要自定义压测程序？

场景 命令行工具局限 自定义库优势
动态请求体 只能读取静态文件或 stdin 运行时生成 JSON/Protobuf 负载
条件停止 固定持续时间或请求数 按 P99 阈值、错误率自动停止
自定义指标 仅输出预定义分位值 采集 goroutine 数、GC 暂停、连接池状态
多阶段压测 需外部编排 代码内串联 warmup → 峰值 → 持续
实战建议： 在搜索尾延迟诊断中，建议始终使用自定义脚本，因为你需要精确控制每次请求的搜索关键词分布、缓存命中策略和超时时间，这些都无法通过纯命令行完成。

4.2 核心依赖与项目结构

引入 vegeta 库（v12 及以上版本）：

go get -u github.com/tsenart/vegeta/v12/lib

推荐的项目文件组织：

latency-bench/
├── main.go              # 入口：解析参数、启动压测
├── target.go            # 动态 Target 构造器
├── metrics.go           # 自定义指标采集与输出
├── config.go            # 压测配置结构体
└── go.mod / go.sum

4.3 基础压测脚本骨架

以下是一个最小可用的自定义压测程序，使用 vegeta.Attacker 控制并发与速率：

package main

import (
    "fmt"
    "time"
    vegeta "github.com/tsenart/vegeta/v12/lib"
)

func main() {
    // 1. 定义压测目标
    target := vegeta.Target{
        Method: "POST",
        URL:    "http://search-api:8080/search",
        Body:   []byte(`{"query":"go latency","page":1}`),
        Header: map[string][]string{
            "Content-Type": {"application/json"},
        },
    }

    // 2. 构造 Targeter（可重复生成目标）
    targeter := vegeta.NewStaticTargeter(target)

    // 3. 创建攻击器：控制并发数
    attacker := vegeta.NewAttacker(
        vegeta.Workers(uint64(10)),       // 并发工作协程数
        vegeta.MaxWorkers(uint64(100)),   // 最大并发上限
        vegeta.Connections(5),            // 每个目标主机的连接数
        vegeta.Timeout(5*time.Second),    // 请求超时
    )

    // 4. 定义速率：每秒 500 次请求
    rate := vegeta.Rate{Freq: 500, Per: time.Second}

    // 5. 执行压测，持续 30 秒
    var metrics vegeta.Metrics
    for res := range attacker.Attack(targeter, rate, 30*time.Second, "search-bench") {
        metrics.Add(res)
    }
    metrics.Close()

    // 6. 输出结果
    fmt.Printf("总请求数: %d\n", metrics.Requests)
    fmt.Printf("成功率: %.2f%%\n", metrics.Success*100)
    fmt.Printf("P50: %s\n", metrics.Latencies.P50)
    fmt.Printf("P99: %s\n", metrics.Latencies.P99)
    fmt.Printf("最大延迟: %s\n", metrics.Latencies.Max)
}
关键参数说明：
  • Workers：初始并发协程数，决定最低并发水位。
  • MaxWorkers：当速率较高时，vegeta 会自动增加协程，此参数限制上限。
  • Connections：每个目标主机的 HTTP 连接复用数，对尾延迟影响显著。
  • Timeout：超过此时间的请求会被强制终止并记录为错误。

4.4 并发与速率控制原理

vegeta 的速率控制基于 令牌桶算法，而非简单的 sleep 间隔。这使得在突发场景下仍能维持平均速率，同时允许短时抖动。

vegeta 并发与速率控制架构 用户配置 速率控制器 (令牌桶) 工作协程池 (Workers) HTTP 连接复用层 目标服务器 指标收集器 (Metrics) 说明：令牌桶控制平均速率，工作协程池控制并发度，连接复用层控制 TCP 连接数，三者独立配置

4.5 动态 Target 构造器

在搜索尾延迟测试中，每次请求的查询词、分页参数、用户 ID 都应不同。使用 vegeta.Targeter 函数签名实现动态生成：

func dynamicTargeter() vegeta.Targeter {
    var counter uint64
    queries := []string{"go latency", "tail latency", "vegeta benchmark", "golang p99"}
    return func(tgt *vegeta.Target) error {
        idx := atomic.AddUint64(&counter, 1) % uint64(len(queries))
        body, _ := json.Marshal(map[string]interface{}{
            "query": queries[idx],
            "page":  int(idx) + 1,
            "user":  fmt.Sprintf("user_%d", counter%1000),
        })
        *tgt = vegeta.Target{
            Method: "POST",
            URL:    "http://search-api:8080/search",
            Body:   body,
            Header: map[string][]string{
                "Content-Type": {"application/json"},
                "X-Request-ID": {fmt.Sprintf("req-%d", counter)},
            },
        }
        return nil
    }
}
注意： Targeter 函数会被多个 goroutine 并发调用，必须保证内部状态（如 counter）的原子操作。推荐使用 sync/atomic 或 sync.Mutex。

4.6 多阶段压测：Warmup + 峰值 + 持续

尾延迟诊断需要观察系统在不同负载阶段的行为。以下代码实现三阶段压测：

func multiStageAttack(targeter vegeta.Targeter) {
    attacker := vegeta.NewAttacker(
        vegeta.Workers(5),
        vegeta.MaxWorkers(50),
        vegeta.Timeout(10*time.Second),
    )

    stages := []struct {
        rate     vegeta.Rate
        duration time.Duration
        name     string
    }{
        {rate: vegeta.Rate{Freq: 100, Per: time.Second}, duration: 30 * time.Second, name: "warmup"},
        {rate: vegeta.Rate{Freq: 1000, Per: time.Second}, duration: 60 * time.Second, name: "peak"},
        {rate: vegeta.Rate{Freq: 500, Per: time.Second}, duration: 120 * time.Second, name: "sustain"},
    }

    for _, stage := range stages {
        fmt.Printf("阶段: %s | 速率: %d req/s | 持续: %s\n",
            stage.name, stage.rate.Freq, stage.duration)

        var metrics vegeta.Metrics
        for res := range attacker.Attack(targeter, stage.rate, stage.duration, stage.name) {
            metrics.Add(res)
        }
        metrics.Close()

        fmt.Printf("  P50: %s | P99: %s | 成功率: %.2f%%\n",
            metrics.Latencies.P50, metrics.Latencies.P99, metrics.Success*100)

        // 阶段间冷却，让系统恢复
        time.Sleep(5 * time.Second)
    }
}

4.7 自定义指标采集：捕获系统级延迟

除了 vegeta 内置的请求延迟，我们还需要采集 goroutine 数量、GC 暂停时间等系统指标，用于关联分析：

func collectSystemMetrics(interval time.Duration, stopCh chan struct{}) {
    ticker := time.NewTicker(interval)
    defer ticker.Stop()

    var memStats runtime.MemStats
    for {
        select {
        case <-ticker.C:
            runtime.ReadMemStats(&memStats)
            log.Printf("[系统指标] goroutines=%d | GC暂停(最近)=%dμs | 堆分配=%dMB",
                runtime.NumGoroutine(),
                memStats.PauseNs[(memStats.NumGC+255)%256]/1000,
                memStats.Alloc/1024/1024)
        case <-stopCh:
            return
        }
    }
}
集成方式： 在 attacker.Attack 启动前，另起一个 goroutine 运行 collectSystemMetrics，并在压测结束后通过 stopCh 关闭。这样可以在压测日志中看到系统状态与延迟的对应关系。

4.8 参数对比速查表

参数 默认值 尾延迟影响 推荐设置（搜索场景）
Workers 10 并发不足导致请求排队，增加 P99 CPU 核心数 × 2
MaxWorkers 无限制 过高导致协程争抢，延迟抖动 Workers × 5
Connections 0（无限制） 连接复用不足增加 TCP 握手延迟 10 ~ 50
Timeout 30s 过长会掩盖超时尾延迟 P99 期望值 × 3
Rate.Freq 50 直接决定系统负载 从 100 开始逐步上调

4.9 完整运行示例

将上述组件组合成一个可执行的压测脚本：

func main() {
    // 1. 动态 targeter
    targeter := dynamicTargeter()

    // 2. 启动系统指标采集
    stopCh := make(chan struct{})
    go collectSystemMetrics(5*time.Second, stopCh)

    // 3. 执行多阶段压测
    multiStageAttack(targeter)

    // 4. 停止指标采集
    close(stopCh)

    fmt.Println("压测完成，请分析尾延迟数据。")
}
输出示例： 
阶段: warmup | 速率: 100 req/s | 持续: 30s
  P50: 12ms | P99: 45ms | 成功率: 100.00%
阶段: peak | 速率: 1000 req/s | 持续: 60s
  P50: 28ms | P99: 210ms | 成功率: 99.87%
阶段: sustain | 速率: 500 req/s | 持续: 120s
  P50: 15ms | P99: 78ms | 成功率: 100.00%
← 上一章 📖 返回目录 下一章 →