go 框架中实现限流保护令牌桶算法是一种限流方法，它可防止系统在高并发场景下过载。在 go 中使用令牌桶算法实现限流 protection：创建令牌桶限流器。尝试获取令牌（允许）。如果有足够的令牌，则获取令牌并允许请求。如果没有令牌，则拒绝请求。
Go 框架中实现限流保护
简介
限流是一个至关重要的技术，可以在高并发场景下保护系统免受过载。在 Go 中实现限流保护有多种方法，本文将探讨使用令牌桶算法实现的一个简单而有效的方案。
立即学习“go语言免费学习笔记（深入）”；
令牌桶算法
令牌桶算法是一种经典的限流算法。它可以想象成一个桶，其中装有令牌。每次请求到达时，就会从桶中取走一个令牌。如果桶中没有令牌，则请求会被拒绝（限流）。令牌不断以固定的速率生成，当请求高峰期过后，桶中的令牌数量会逐渐恢复。
实现
在 Go 中实现令牌桶算法很简单：
import (
"context"
"sync"
"time"
)
// 令牌桶限流器
type TokenBucketLimiter struct {
// 最大令牌数
maxTokens int
// 当前令牌数
tokens int
// 令牌生成速率（令牌数/秒）
rate float64
// 上次生成令牌时间
lastGen time.Time
// 互斥锁
lock sync.Mutex
}
// NewTokenBucketLimiter 创建新的令牌桶限流器
func NewTokenBucketLimiter(maxTokens int, rate float64) *TokenBucketLimiter {
return &TokenBucketLimiter{
maxTokens: maxTokens,
tokens: 0,
rate: rate,
lastGen: time.Now(),
}
}
// Allow 尝试获取令牌
func (l *TokenBucketLimiter) Allow(ctx context.Context) bool {
l.lock.Lock()
defer l.lock.Unlock()
// 计算当前令牌数
now := time.Now()
l.tokens += int(l.rate * now.Sub(l.lastGen).Seconds())
if l.tokens > l.maxTokens {
l.tokens = l.maxTokens
}
l.lastGen = now
// 如果有足够的令牌，则获取令牌并允许请求
if l.tokens > 0 {
l.tokens--
return true
}
// 如果没有令牌，则拒绝请求
return false
}实战案例
下面是一个简单的示例，展示如何使用令牌桶限流器限制对 HTTP 路由的请求：
import (
"log"
"net/http"
"github.com/go-chi/chi/v5/middleware"
"your-app/ratelimit"
)
// 创建令牌桶限流器
limiter := ratelimit.NewTokenBucketLimiter(100, 10)
// 创建 HTTP 路由
r := chi.NewRouter()
// 使用限流器中间件保护路由
r.Use(middleware.Limit(limiter))
// 处理请求
r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
if limiter.Allow(r.Context()) {
// 允许请求，执行业务逻辑
} else {
// 拒绝请求，返回错误
http.Error(w, "Service is busy, please try again later.", http.StatusTooManyRequests)
}
})结论
使用令牌桶算法实现限流保护是简单且有效的，它可以防止系统在高并发场景下过载。本文介绍的 Go 框架实现提供了易于使用的 API，可以轻松集成到您的应用程序中。在实现限流时，请仔细选择令牌桶的参数（最大令牌数和生成速率），以获得最佳效果。
