接口被刷爆？用 Go+Gin 套一层令牌桶限流

一、几种常见的限流算法

限流这事说简单也简单，就是控制单位时间内放过去多少请求。但具体怎么控，有好几种思路，我列个表对比下：

我选令牌桶，原因就一条：它能容忍短时突发。大部分业务接口都有这个特征——平时流量平稳，偶尔来一波高峰，你不想把正常用户也挡在外面。漏桶太死板，固定窗口有突刺问题，滑动窗口实现起来啰嗦。令牌桶刚好卡在灵活和简单之间。

二、令牌桶到底怎么转的

画了个图，比干看文字直观点：

核心就几个参数：

• 桶容量（capacity）：桶里最多放多少令牌，也是允许的最大突发量
• 发放速率（rate）：每秒往桶里加多少令牌
• 当前令牌数（tokens）：桶里现在有多少令牌，随时间增长

一个请求来了，先看桶里有没有令牌。有就拿走一个，放行；没有就拒绝。同时后台有个逻辑：每过一段时间往桶里补令牌，但不超过桶容量。

这里有个偷懒的技巧——补令牌不需要真的起个 goroutine 定时往里加，算个时间差就行。上次请求到现在过了多久，该补多少令牌，一减一加就出来了。这个惰性计算的思路省了一个定时器，代码干净，还少一层并发问题。

三、中间件实现

先定义桶结构：

package limiter

import (
	"sync"
	"time"
)

// TokenBucket 令牌桶
type TokenBucket struct {
	rate       float64   // 每秒发放的令牌数
	capacity   float64   // 桶容量
	tokens     float64   // 当前令牌数
	lastRefill time.Time // 上次补充时间
	mu         sync.Mutex
}

// NewTokenBucket 创建一个令牌桶
func NewTokenBucket(rate, capacity float64) *TokenBucket {
	return &TokenBucket{
		rate:       rate,
		capacity:   capacity,
		tokens:     capacity, // 初始满桶
		lastRefill: time.Now(),
	}
}

然后是核心的 Allow 方法——拿令牌：

// Allow 尝试获取一个令牌，成功返回 true
func (tb *TokenBucket) Allow() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	now := time.Now()
	// 算算从上次到现在该补多少令牌
	elapsed := now.Sub(tb.lastRefill).Seconds()
	tb.tokens += elapsed * tb.rate

	// 别超过桶容量
	if tb.tokens > tb.capacity {
		tb.tokens = tb.capacity
	}
	tb.lastRefill = now

	// 有令牌就拿走一个
	if tb.tokens >= 1 {
		tb.tokens -= 1
		return true
	}
	return false
}

这段代码不长，但有几个点值得说道。elapsed * tb.rate 就是这段时间该补的令牌数，乘出来是个浮点数，没关系，令牌数用 float64 存就行，判断的时候看够不够 1 个就成。加锁用 sync.Mutex 就够了，限流中间件不是热点路径，锁竞争不严重。要是你实在嫌锁慢，可以换 sync/atomic 做 CAS，但代码复杂度上去了，个人觉得没必要。

接下来包成 Gin 中间件：

package middleware

import (
	"net/http"
	"sync"

	"github.com/gin-gonic/gin"
	"yourproject/limiter"
)

// IPRateLimiter 按客户端IP分别建桶
type IPRateLimiter struct {
	limiters  map[string]*limiter.TokenBucket
	mu        sync.Mutex
	rate      float64
	capacity  float64
}

func NewIPRateLimiter(rate, capacity float64) *IPRateLimiter {
	return &IPRateLimiter{
		limiters: make(map[string]*limiter.TokenBucket),
		rate:     rate,
		capacity: capacity,
	}
}

// getLimiter 获取某个IP的限流器，没有就建一个
func (i *IPRateLimiter) getLimiter(ip string) *limiter.TokenBucket {
	i.mu.Lock()
	defer i.mu.Unlock()

	l, exists := i.limiters[ip]
	if !exists {
		l = limiter.NewTokenBucket(i.rate, i.capacity)
		i.limiters[ip] = l
	}
	return l
}

// RateLimit 限流中间件
func (i *IPRateLimiter) RateLimit() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		l := i.getLimiter(ip)

		if !l.Allow() {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"code":    429,
				"message": "请求太频繁了，稍后再试试",
			})
			c.Abort()
			return
		}
		c.Next()
	}
}

这里按 IP 分别建桶，每个 IP 独立限流。你也可以改成按用户 ID、按接口路径来分，逻辑一样，换个 key 就行。

四、接入 Gin 路由

写个 main 函数把中间件挂上去：

package main

import (
	"log"

	"github.com/gin-gonic/gin"
	"yourproject/middleware"
)

func main() {
	r := gin.Default()

	// 每秒10个令牌，桶容量20（允许短时突发到20）
	limiter := middleware.NewIPRateLimiter(10, 20)

	// 给所有路由加限流
	r.Use(limiter.RateLimit())

	r.GET("/api/order", func(c *gin.Context) {
		c.JSON(200, gin.H{"msg": "订单查询成功"})
	})

	r.GET("/api/user", func(c *gin.Context) {
		c.JSON(200, gin.H{"msg": "用户信息"})
	})

	log.Fatal(r.Run(":8080"))
}

跑起来就 go run main.go，没别的花活。拿 curl 连着打几下就能看到效果——前 20 个请求返回 200，之后的开始返回 429，过一秒又能打 10 个。

五、压测看看效果

我用 hey 压了一波，对比开限流前后的表现：

开了限流之后，大部分超量请求直接 429 打回去了，根本不进业务逻辑，所以成功的那些请求延迟反而更低了。被拒绝的 820 个请求几乎是瞬时返回的，不占数据库连接，不占 goroutine，很轻。限流的价值就在这——把流量削到系统能扛得住的水平，该放行的放行，扛不住的打回去。

六、上线前要注意的几点

桶大小怎么定？ 我的经验是先看正常峰值 QPS，桶容量设到峰值的 2-3 倍，发放速率设到峰值水平。比如平时峰值 50 QPS，桶容量设 100 到 150，rate 设 50。上线后看监控再微调，别想一步到位。

限流粒度选哪个？ 按 IP 限流最简单，但 NAT 后面一堆用户共用一个 IP，容易误伤。如果是 to C 的服务，建议按用户 ID 限流，登录态接口天然有 user_id。如果是公开 API，按 IP 限流就够用了，别搞太复杂。

map 会无限增长吗？ 上面代码里 limiters 这个 map，每个新 IP 都会加一条，理论上会一直涨。生产环境得加个过期清理。简单做法是起个定时器，每分钟扫一遍，把超过 5 分钟没请求的 IP 删掉：

func (i *IPRateLimiter) Cleanup(interval, maxIdle time.Duration) {
	ticker := time.NewTicker(interval)
	go func() {
		for range ticker.C {
			i.mu.Lock()
			for ip, l := range i.limiters {
				if time.Since(l.LastRefill()) > maxIdle {
					delete(i.limiters, ip)
				}
			}
			i.mu.Unlock()
		}
	}()
}

lastRefill 字段是私有的，得加个 LastRefill() 方法暴露出去，这个小改我就不贴了。

分布式环境怎么办？ 上面这套是单机限流，如果你有 3 台机器负载均衡，每台各限 10 QPS，加起来就是 30 QPS。要是你需要全局限流 10 QPS，得把令牌桶状态放到 Redis 里。用 Redis 的 INCR 加 EXPIRE 能搞个简化版，或者直接用 redis-cell 这个 Redis 模块，它原生支持令牌桶算法。不过单机限流对大部分中小项目够用了，别一上来就过度设计。

总结一下

这个令牌桶中间件从写到上线大概一个半小时，代码不到 100 行，但效果实实在在。限流这事跟加锁一样，属于那种"不出事没人想加，出了事后悔没加"的东西。。。