go gin框架实现限流
接手一个项目,甲方希望在项目中实现指定用户的限流,这需要我们在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) } }
打完收工!