17370845950

本文介绍使用 `juju/ratelimit` 库结合令牌桶算法，为 http 文件上传和下载添加可配置的速率限制（如 1mb/s），适用于大文件（如 1gb）场景，避免带宽耗尽或服务阻塞。

在 Go 编写的 HTTP 文件服务中，若不对上传/下载流量进行节流，单个大文件传输可能占满服务器带宽，影响其他请求响应甚至导致超时。Go 标准库本身不提供内置限速机制，但可通过包装 io.Reader 和 io.Writer 的方式，在数据流关键路径上插入速率控制器。推荐使用成熟、轻量且线程安全的第三方库 github.com/juju/ratelimit，其基于经典的令牌桶（Token Bucket）算法实现：以恒定速率向桶中注入令牌，每次读/写操作需消耗令牌，从而自然实现平滑限速。

✅ 限速原理简述

• NewBucketWithRate(fillRate, capacity) 创建一个桶：
  • fillRate：每秒新增令牌数（单位：字节/秒），即目标限速值（如 1 * 1024 * 1024 表示 1MB/s）；
  • capacity：桶最大容量（建议 ≥ fillRate，避免初始突发流量被过度抑制）。
• ratelimit.Reader(r, bucket) 返回一个限速包装的 io.Reader，每次 Read() 前自动等待足够令牌；
• 同理，ratelimit.Writer(w, bucket) 可用于限速写入（适用于下载响应流）。

?️ 实际集成示例

1. 上传限速（修改 uploadFile）

import (
    "io"
    "net/http"
    "os"
    "github.com/juju/ratelimit"
)

func uploadFile(w http.ResponseWriter, r *http.Request) {
    file, _, err := r.FormFile("file")
    if err != nil {
        http.Error(w, "无法获取上传文件: "+err.Error(), http.StatusBadRequest)
        return
    }
    defer file.Close()

    // 创建限速桶：1MB/s，桶容量 1MB（允许短时突发）
    bucket := ratelimit.NewBucketWithRate(1*1024*1024, 1*1024*1024)

    // 包装上传文件流，实现上传速率限制
    limitedReader := ratelimit.Reader(file, bucket)

    os.MkdirAll(`e:\test`, 0755)
    out, err := os.Create(`e:\test\test.mpg`)
    if err != nil {
        http.Error(w, "无法创建目标文件: "+err.Error(), http.StatusInternalServerError)
        return
    }
    defer out.Close()

    // 使用限速 Reader 进行拷贝
    _, err = io.Copy(out, limitedReader)
    if err != nil {
        http.Error(w, "文件写入失败: "+err.Error(), http.StatusInternalServerError)
        return
    }

    w.WriteHeader(http.StatusOK)
    w.Write([]byte("上传完成，已限速"))
}

2. 下载限速（添加 downloadFile）

func downloadFile(w http.ResponseWriter, r *http.Request) {
    f, err := os.Open(`e:\test\test.mpg`)
    if err != nil {
        http.Error(w, "文件不存在", http.StatusNotFound)
        return
    }
    defer f.Close()

    // 设置响应头
    w.Header().Set("Content-Type", "application/octet-stream")
    w.Header().Set("Content-Disposition", `attachment; filename="test.mpg"`)

    // 创建限速桶（同样 1MB/s）
    bucket := ratelimit.NewBucketWithRate(1*1024*1024, 1*1024*1024)

    // 包装 ResponseWriter 为限速 Writer
    limitedWriter := ratelimit.Writer(w, bucket)

    // 流式传输，自动限速
    _, err = io.Copy(limitedWriter, f)
    if err != nil && err != io.ErrUnexpectedEOF {
        // 忽略客户端提前断开连接（常见于大文件下载中断）
        log.Printf("下载时发生错误: %v", err)
    }
}

⚠️ 注意事项与最佳实践

• 桶容量选择：capacity 过小会导致频繁等待，增大延迟；过大则削弱限速效果。建议设为 fillRate 的 1–3 倍，平衡平滑性与响应性。
• 并发安全性：ratelimit.Bucket 是并发安全的，可复用同一个桶控制多个连接（如全局限速），也可为每个请求新建桶实现独立限速（推荐用于用户级配额）。
• 错误处理：io.Copy 在限速下仍可能因网络中断返回 io.ErrUnexpectedEOF，应忽略该错误（客户端主动断连），仅记录非预期错误。
• 资源释放：务必确保 file.Close() 和 out.Close() 被调用，避免句柄泄漏；限速包装器不改变底层资源生命周期。
• 配置化扩展：将速率值（如 1 * 1024 * 1024）提取为变量或从配置文件/环境变量加载，便于动态调整。

通过上述方式，你可以在不修改核心业务逻辑的前提下，精准、低侵入地为 Go Web 服务添加生产级的上传/下载速率控制能力。