golang gin框架中实现一个简单的不是特别精确的秒级限流器

起因

看了两篇关于golang中限流器的帖子：

• Gin 开发实践：如何实现限流中间件
• 常用限流策略——漏桶与令牌桶介绍

我照着用，居然没效果……

时间有限没有深究。这实在是一个很简单的功能，我的需求是每分钟限制大约xxx次请求，平均每秒限制到XXX除60次左右的请求也就够了。

限流器代码

于是下面用了个及其简单的方式来实现限流器：

utils/limiter.go

package utils

import (
	"sync/atomic"
	"time"
)

//Limiter 限流器对象
type Limiter struct {
	value  int64
	max    int64
	ts     int64
}

//NewLimiter 产生一个限流器
func NewLimiter(cnt int64) *Limiter {
	return &Limiter{
		value:  0,
		max:    cnt,
		ts:     time.Now().Unix(),
	}
}

//Ok 是否可以通过
func (l *Limiter) Ok() bool {
	ts := time.Now().Unix()
	tsOld := atomic.LoadInt64(&l.ts)
	if ts != tsOld {
		atomic.StoreInt64(&l.ts, ts)
		atomic.StoreInt64(&l.value, 1)
		return true
	}
	return atomic.AddInt64(&(l.value), 1) < l.max
}

//SetMax 设置最大限制
func (l *Limiter) SetMax(m int64) {
	l.max = m
}

gin框架中使用

main.go

//MaxAllowed 限流器
func MaxAllowed(limitValue int64) func(c *gin.Context) {
	limiter := utils.NewLimiter(limitValue)
	log.Println("limiter.SetMax:", limitValue)
	// 返回限流逻辑
	return func(c *gin.Context) {
		if !limiter.Ok() {
			c.AbortWithStatus(http.StatusServiceUnavailable)  //超过每秒200，就返回503错误码
			return
		}
		c.Next()
	}
}

func main(){
    router := gin.New()
    router.Use(MaxAllowed(200))  //限制每秒最多允许200个请求
}

have fun!