Go实现限流器

Tags:

在 Go 中实现一个简单的 Rate Limiter（限流器） 可以有多种方式，常见的有基于 令牌桶（Token Bucket） 和 漏桶（Leaky Bucket） 的实现。下面我们先介绍一种经典方式 —— 令牌桶算法 实现。

1. 令牌桶算法

“令牌桶算法（Token Bucket）”是一种 限流算法，用于控制请求的速率，是常见于 API 网关、负载均衡器、流控中间件等的核心策略之一。

1.1. 核心思想

系统以固定速率往“桶”里放令牌，每次请求必须拿到一个令牌才能被处理，否则被拒绝（限流）。

名称 含义
桶（Bucket） 保存令牌的容器，最多可以装 burst 个令牌（突发容量）
令牌（Token） 每个令牌代表一次允许的操作（如一次 HTTP 请求）
速率（Rate） 向桶中添加令牌的速度，比如每秒放 5 个
请求到来时 若桶中有令牌，请求就“取出一个令牌”被允许；否则就被拒绝或等待

1.2. 动态行为图示

假设：每秒生成1个令牌，桶最多装3个。

时间：0s 桶：[●] 请求：允许（消耗1个令牌）
时间：1s 桶：[●] 请求：允许（消耗1个令牌）
时间：2s 桶：[●●] 请求：允许
时间：3s 桶：[●●●] 请求：允许
时间：4s 桶：[●●●] 请求：没有消费，桶满了（不再增加）
时间：5s 桶：[●●●] 来5个请求，只允许3个，其余被限流

1.3. 和漏桶算法的对比

比较项 令牌桶（Token Bucket） 漏桶（Leaky Bucket）
控制方式 控制请求进入速率 控制请求处理速率
支持突发请求 ✅ 支持 ❌ 严格匀速处理
应用场景 限流（API、接口） 排队（网络、处理任务）

1.4. 应用场景举例

  • API 请求速率控制（每个用户最多1秒10次）

  • 登录/注册防暴力攻击（每 IP 限1分钟5次）

  • CDN 边缘限流（防止源站被打爆）

  • 后端任务投递（防止 RabbitMQ 拥堵）

2. 实现简单 Token Bucket 限流器

2.1. 使用 time.Ticker 实现

package main

import (
    "fmt"
    "time"
)

type RateLimiter struct {
    tokens       chan struct{}
    ticker       *time.Ticker
    maxTokens    int
    refillPeriod time.Duration
}

func NewRateLimiter(rps int, burst int) *RateLimiter {
    rl := &RateLimiter{
        tokens:       make(chan struct{}, burst),
        ticker:       time.NewTicker(time.Second / time.Duration(rps)),
        maxTokens:    burst,
        refillPeriod: time.Second / time.Duration(rps),
    }

    // 初始化 token 桶
    for i := 0; i < burst; i++ {
        rl.tokens <- struct{}{}
    }

    // 后台协程定时放 token
    go func() {
        for range rl.ticker.C {
            select {
            case rl.tokens <- struct{}{}:
            default: // 桶满时丢弃 token
            }
        }
    }()

    return rl
}

// Allow 会阻塞直到有 token 可用（可改成 TryAllow 非阻塞版本）
func (rl *RateLimiter) Allow() bool {
    select {
    case <-rl.tokens:
        return true
    default:
        return false
    }
}

func main() {
    limiter := NewRateLimiter(5, 10) // 每秒 5 个请求，最多缓冲 10 个

    for i := 0; i < 20; i++ {
        if limiter.Allow() {
            fmt.Println("Request", i, "allowed at", time.Now())
        } else {
            fmt.Println("Request", i, "denied at", time.Now())
        }
        time.Sleep(100 * time.Millisecond)
    }
}

2.2. 使用Go rate包实现

Go 的 golang.org/x/time/rate 包就实现了 令牌桶算法，例如：

// 每秒5个令牌，最多能装10个（允许突发10次）
limiter := rate.NewLimiter(5, 10)

你可以使用 .Allow() 来判断是否获得令牌：

if limiter.Allow() {
    // 有令牌，请求通过
} else {
    // 没令牌，请求被限流
}

3. 实现Gin按用户或 IP 限流中间件

3.1. 限流器结构（使用 rate.Limiter）

package ratelimiter

import (
    "sync"

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

type RateLimiter struct {
    rate     rate.Limit
    burst    int
    buckets  map[string]*rate.Limiter
    mutex    sync.Mutex
}

func NewRateLimiter(r rate.Limit, b int) *RateLimiter {
    return &RateLimiter{
        rate:    r,
        burst:   b,
        buckets: make(map[string]*rate.Limiter),
    }
}

func (rl *RateLimiter) getLimiter(key string) *rate.Limiter {
    rl.mutex.Lock()
    defer rl.mutex.Unlock()
    if limiter, exists := rl.buckets[key]; exists {
        return limiter
    }
    limiter := rate.NewLimiter(rl.rate, rl.burst)
    rl.buckets[key] = limiter
    return limiter
}

func (rl *RateLimiter) Allow(key string) bool {
    return rl.getLimiter(key).Allow()
}

3.2. Gin 中间件

package ratelimiter

import (
    "net/http"

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

type KeyFunc func(c *gin.Context) string

func Middleware(rl *RateLimiter, keyFn KeyFunc) gin.HandlerFunc {
    return func(c *gin.Context) {
        key := keyFn(c)
        if !rl.Allow(key) {
            c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
                "error": "rate limit exceeded",
            })
            return
        }
        c.Next()
    }
}

3.3. 在 main.go 中使用（按 IP 限流）

package main

import (
    "github.com/gin-gonic/gin"
    "golang.org/x/time/rate"

    "your_project/ratelimiter"
)

func main() {
    // 每秒2个请求，最多突发5个
    rl := ratelimiter.NewRateLimiter(2, 5)

    r := gin.Default()

    r.Use(ratelimiter.Middleware(rl, func(c *gin.Context) string {
        return c.ClientIP() // 或使用 c.GetString("userID") 实现按用户限流
    }))

    r.GET("/hello", func(c *gin.Context) {
        c.JSON(200, gin.H{"msg": "hello"})
    })

    r.Run(":8080")
}

最后修改 June 6, 2025: go ratelimiter (2046a97)