Go从零实现高性能Token Bucket限流器：应对API突发流量与脚本攻击

你是否还记得自己开发的API第一次爆火时的场景？短暂的喜悦之后，可能就是无尽的告警。曾经，一个疑似脚本的用户就把我的服务打到了近乎停滞，那种精心打磨的快速API瞬间变成“粘稠API”的挫败感，至今难忘。

那次经历给我上了深刻的一课：服务火起来是好事，但流量失控就成了灾难。无脑地增加服务器并非良策，成本飙升不说，还可能治标不治本。很多时候，你需要的不是一个庞大的服务团队，而是一个懂得看菜下碟的智能门卫——也就是我们今天要讨论的限流器。

本文将带你用 Go 从零实现一个我最钟爱的限流算法：Token Bucket（令牌桶）。用 Go 来实现它，可以说非常契合这门语言的气质。

限流究竟有何意义？🤔

在动手写代码之前，我们先对齐一下认知：限流到底在解决什么问题？

想象一个火爆的夜店，消防规定最多只能容纳 100 人。门口保安的作用就是确保这个规则被执行，即使一下子来了50个人，也不能全部放行，必须等里面有人出来才行。这个保安，就是Rate Limiter（速率限制器）。

在 API 的世界里，逻辑完全一致。它的重要性至少体现在以下几个方面：

✅ 防止服务器过载：当某个用户过于“热情”时，限流能保护其他用户的体验依然丝滑。
✅ 增强安全性：应对暴力破解密码、刷接口、撞库等攻击，限流是第一道防线。
✅ 保证公平性：避免单一用户或客户端独占所有资源，让其他请求“无饭可吃”。
✅ 控制成本：在云服务按请求计费的时代，有效的限流能直接帮你减少账单上的“惊喜”。

说真的，一个对外开放的 API 如果不做限流，无异于敞开大门并挂上“内有珍宝，免费自取”的牌子。

核心算法：Token Bucket（令牌桶）🪣

那么这个算法是如何工作的呢？名字听起来有点玄乎，但原理非常直观。想象一个桶。

核心规则：每个想要进来的请求，都必须持有一枚令牌（Token）。没有令牌？抱歉，请稍后再试。

我们可以将它的运行机制分解如下：

• 桶（Bucket）：容量固定，比如最多能装 100 个令牌。装满后，再往里添加也会溢出。
• 令牌（Token）：一枚令牌相当于一次请求的通行证。
• 补充（Refill）：系统以一个固定的速率向桶中添加令牌，例如“每秒添加 10 个”。
• 请求（Request）：当请求到达时，它会尝试从桶中取走一枚令牌：
  • 有令牌：成功取走，请求被放行 ✅
  • 无令牌：获取失败，请求被拒绝 ❌

我特别喜欢 Token Bucket 算法的一点在于：它对突发流量（Burst Traffic）非常友好。当桶是满的时候，你可以一瞬间放行大量请求（例如100个），享受短暂的爆发处理能力；而在爆发之后，请求速率又会被令牌补充速率严格限制，回归稳态。它不像某些简单粗暴的方案，而是做到了既灵活，又可控。

开始用 Go 实现！💻

好了，理论铺垫完毕，开始动手。我们的目标是实现一个简单且线程安全的 Token Bucket 限流器。

这里有一个关键设计点：我打算使用 Go 语言内置的 channel 来充当这个“桶”。原因非常实际：

• Channel 天生就是并发安全的。
• 使用得当可以避免显式锁（Mutex）带来的复杂度。
• 配合 select 语句可以实现非阻塞操作，写出来的代码既快又优雅。

首先，定义我们的结构体：

package main

import (
    "fmt"
    "sync"
    "time"
)

// TokenBucket 保存限流器的所有状态。
type TokenBucket struct {
    capacity    int
    tokens      chan struct{} // 这个 channel 就是桶本体
    refillEvery time.Duration
    stop        chan struct{}
}

// NewTokenBucket 构造函数：创建桶并启动补充协程。
func NewTokenBucket(capacity int, refillEvery time.Duration) *TokenBucket {
    tb := &TokenBucket{
        capacity:    capacity,
        tokens:      make(chan struct{}, capacity),
        refillEvery: refillEvery,
        stop:        make(chan struct{}),
    }

    // 我更喜欢让桶“初始状态就是满的”，这样它的突发处理能力更符合直觉（也贴合多数生产需求）
    for i := 0; i < capacity; i++ {
        tb.tokens <- struct{}{}
    }

    go tb.refillLoop()
    return tb
}

这里我做了一个关键的优化：初始化时直接将桶填满。很多实际系统希望服务在启动后就具备处理突发流量的能力。否则，刚重启完毕，所有请求一上来就被拒绝，用户体验会非常糟糕。

令牌补充协程（后台“滴水”）

// refillLoop 在单独的 goroutine 里周期性补充令牌。
func (tb *TokenBucket) refillLoop() {
    ticker := time.NewTicker(tb.refillEvery)
    defer ticker.Stop()

    for {
        select {
        case <-ticker.C:
            // 尝试补充一个令牌。如果桶满了就直接跳过，绝不阻塞。
            select {
            case tb.tokens <- struct{}{}:
            default:
                // 桶已满，放弃本次补充，严格维持上限
            }
        case <-tb.stop:
            return
        }
    }
}

这里的精髓在于 select 语句配合 default 分支：

• 向已满的 channel 写入数据会导致阻塞。
• 而 default 分支让我们能够实现“写不进去就算了”的逻辑，不会拖慢整个系统。
  这正是高性能限流器应有的“干脆利落”。

核心门禁函数：Allow

// Allow 尝试获取一个令牌：拿到就放行，拿不到就拒绝。
// 非阻塞、速度极快。
func (tb *TokenBucket) Allow() bool {
    select {
    case <-tb.tokens:
        return true
    default:
        return false
    }
}

// Stop 停止后台的令牌补充协程（用于优雅退出）。
func (tb *TokenBucket) Stop() {
    close(tb.stop)
}

Allow 方法的逻辑非常“保安”：

• 口袋里还有票（Token）吗？有就给你一张，进去吧。
• 没票了？不好意思，请在门口等候。
  这种实现在高并发场景下非常稳定，因为 channel 操作本身就是并发安全的，而且整个过程无锁、非阻塞，性能很好。

运行演示 🧩

让我们写一个简单的 Demo 来测试它：设置桶容量为 5，每 200 毫秒补充 1 个令牌。这相当于 稳态速率 5 次请求/秒，并允许 瞬时突发 5 个请求。

func main() {
    // 桶最多5个token，每200ms补1个 => 5 req/s
    limiter := NewTokenBucket(5, 200*time.Millisecond)
    defer limiter.Stop()

    fmt.Println("来吧…同时发起 10 个请求！")

    var wg sync.WaitGroup

    for i := 0; i < 10; i++ {
        wg.Add(1)
        go func(reqNum int) {
            defer wg.Done()

            if limiter.Allow() {
                fmt.Printf("✅ Request #%d 放行\n", reqNum)
            } else {
                fmt.Printf("❌ Request #%d 拒绝\n", reqNum)
            }
        }(i + 1)
    }

    wg.Wait()

    fmt.Println("\n歇两秒，让桶回血…")
    time.Sleep(2 * time.Second)

    fmt.Println("再来 5 个请求：")

    for i := 0; i < 5; i++ {
        wg.Add(1)
        go func(reqNum int) {
            defer wg.Done()

            if limiter.Allow() {
                fmt.Printf("✅ Request #%%d 放行\n", reqNum)
            } else {
                fmt.Printf("❌ Request #%%d 拒绝\n", reqNum)
            }
        }(i + 11)
    }

    wg.Wait()
}

运行这段代码，你会看到：

• 第一波 10 个并发请求中，大约前 5 个会被立即放行，后面的则被拒绝。
• 等待 2 秒让令牌桶补充后，接下来的 5 个请求基本都能通过。

其行为完全符合预期：先消耗存量（Burst），再依赖增量（Rate）。

生产环境的选择 📌

我们手动实现的过程很有启发性，但在实际生产环境中，你很可能直接使用 Go 官方维护的优秀库：

import "golang.org/x/time/rate"

limiter := rate.NewLimiter(rate.Limit(5), 5) // 5 req/s, burst=5
if limiter.Allow() {
    // 放行
}

x/time/rate 包更加健壮，测试覆盖更完整，对各种边界情况的处理也更细致。在构建需要应对大规模并发的分布式系统时，这类经过充分验证的组件通常是更可靠的选择。

然而，这绝不意味着我们的“手撸”没有价值。只有自己亲手实现过一次，你在使用这些高级库时才不会“迷信”，而是真正地“驾驭”。理解底层原理的人，在排查限流相关问题时速度会快上一倍；在设计系统架构时，也更不容易被复杂的配置参数所迷惑。

希望这篇从零实现的指南，能帮助你更好地理解 Token Bucket 这一经典的算法思想，并能在你的 Go 项目中自如地运用限流技术，构建出更稳健的服务。欢迎在 云栈社区 分享你的实践与见解。