接手一个项目，甲方希望在项目中实现指定用户的限流，这需要我们在go项目中实现一个限流器。
限流器的实现有很多经典的思想，不过令牌桶的思路简单，运行效率高，它是以恒定的速度往木桶里加入令牌，木桶满了则不再加入令牌。
服务收到请求时尝试从木桶中取出一个令牌，如果能够得到令牌则继续执行后续的业务逻辑。如果没有得到令牌，直接返回访问频率超限的错误码或页面等，不继续执行后续的业务逻辑。
特点：由于木桶内只要有令牌，请求就可以被处理，所以令牌桶算法可以支持突发流量。
同时由于往木桶添加令牌的速度是恒定的，且木桶的容量有上限，所以单位时间内处理的请求书也能够得到控制，起到限流的目的。
假设加入令牌的速度为 1token/10ms，桶的容量为500，在请求比较的少的时候（小于每10毫秒1个请求）时，木桶可以先"攒"一些令牌（最多500个）。
当有突发流量时，一下把木桶内的令牌取空，也就是有500个在并发执行的业务逻辑，之后要等每10ms补充一个新的令牌才能接收一个新的请求。
不过官方已经写好了令牌桶的限流器，我们直接拿过来写好中间件放到路由就好~
middleware/rateLimit.go
package middleware
import (
"IpDatabase/model"
"sync"
"github.com/gin-gonic/gin"
"golang.org/x/time/rate"
)
var limiterMap = make(map[string]*rate.Limiter)
var mu sync.Mutex
// getLimiter 获取或创建一个特定 key 的限流器
func getLimiter(key string) *rate.Limiter {
mu.Lock()
defer mu.Unlock()
// 如果限流器已经存在，返回现有的限流器
if limiter, exists := limiterMap[key]; exists {
return limiter
}
// 否则，创建一个新的限流器
key_limit_info, err := model.ListUser(&model.UserInfo{
Api_key: key,
})
if err != nil {
return nil
}
limiter := rate.NewLimiter(rate.Limit(key_limit_info[0].ReqRate), key_limit_info[0].ReqRate) // 每秒1个请求，突发3个请求
limiterMap[key] = limiter
return limiter
}
// 限流中间件，基于用户的 key（如API Key或User ID）
func RateLimiterByKey(c *gin.Context) {
// 假设用户通过 `Authorization` header 提供 key
key := c.GetHeader("Authorization")
// 如果没有提供 key，返回错误
if key == "" {
BadResponse("API Key is required", c)
return
}
// 获取或创建限流器
limiter := getLimiter(key)
// 检查是否允许请求
if !limiter.Allow() {
BadResponse("Too many requests", c)
return
}
// 如果通过限流器，继续处理请求
c.Next()
}写进路由
func setApiRouter(engine *gin.Engine) {
apiRouter := engine.Group("/api")
{
apiRouter.GET("/search_ip_info", middleware.IsExpiredKey(), middleware.RateLimiterByKey, controller.SearchIp)
}
}打完收工！
文章评论