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
}
//中间件
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))