go语言使用令牌桶限流

安装第三方库

go get -u github.com/juju/ratelimit

代码编写

//limter.go
type LimterIface interface {

	Key(c *gin.Context)string
	GetBucket(key string) (* ratelimit.Bucket,bool)
	AddBuckets(rules ...LimiterBucketRule) LimterIface

}
type Limter struct {
	limterBuckets map[string]*ratelimit.Bucket
}
type LimiterBucketRule struct {
	Key string
	FillInterval time.Duration
	Capacity int64
	Quantum int64
}

//method-limter.go
type MethodLimiter struct {
	*Limter
}

func NewMethodLimiter()LimterIface  {
	return MethodLimiter{
		Limter:&Limter{limterBuckets: make(map[string]*ral4terlimit.Bucket)},
	}
}
func (l MethodLimiter)Key(c *gin.Context)string  {
	uri:=c.Request.RequestURI
	index := strings.Index(uri, "?")
	if index==-1{
		return uri
	}
	return uri[:index]
}
func (l MethodLimiter)GetBucket(key string)(*ratelimit.Bucket,bool){
	bucket,ok := l.limterBuckets[key]
	return bucket,ok
}
func (l MethodLimiter)AddBuckets(rules ...LimiterBucketRule)LimterIface{
	for _,rule:=range rules{
		if _,ok:=l.limterBuckets[rule.Key];!ok{
			ratelimit.NewBucketWithQuantum(rule.FillInterval,rule.Capacity,rule.Quantum)
		}
	}
	return l
}

gin中间件调用

gin中间件调用
//中间件
func Ratelimier(l limter.LimterIface)gin.HandlerFunc{
	return func(c *gin.Context) {
		key:=l.Key(c)
		if bucket,ok:=l.GetBucket(key);ok{
			count:=bucket.TakeAvailable(1)
			if count==0{
				reponse:=app.NewResponse(c)
				reponse.ToErrorReponse(errcode.TooManyRequests)
				c.Abort()
				return
			}
		}
		c.Next()
	}
}
//中间件调用
var methodLimiters=limter.NewMethodLimiter().AddBuckets(limter.LimiterBucketRule{
		Key:          "/auth",
		FillInterval: time.Second,
		Capacity:     10,
		Quantum:      10,
    })
r.Use(middleware.Ratelimier(methodLimiters))

控制请求时间代码

//中间件
func ContextTimeOut(t time.Duration)func(c *gin.Context){
	return func(c *gin.Context) {
		ctx, cance := context.WithTimeout(c.Request.Context(), t)
	  c.Request=c.Request.WithContext(ctx)
	  c.Next()
	  defer cance()
	}
}
//gin中间件调用
r.Use(middleware.ContextTimeOut(time.Minute))