|
|
// 对API接口的QPS进行限制
|
|
|
package traffic
|
|
|
|
|
|
import (
|
|
|
"sync"
|
|
|
"time"
|
|
|
|
|
|
"github.com/mafanr/juz/misc"
|
|
|
)
|
|
|
|
|
|
var apiRates = &sync.Map{}
|
|
|
|
|
|
type apiRate struct {
|
|
|
concurrent int
|
|
|
conUpdate time.Time
|
|
|
|
|
|
qps int
|
|
|
qpsUpdate time.Time
|
|
|
|
|
|
params *sync.Map
|
|
|
}
|
|
|
|
|
|
type param struct {
|
|
|
times int
|
|
|
lastTime time.Time
|
|
|
}
|
|
|
|
|
|
// 能调用此函数，说明qps、并发至少有一项受到了限制
|
|
|
func (rl *RateLimiter) IncApiRate(req misc.TrafficConReq, reply *misc.TrafficConRes) error {
|
|
|
reply.Suc = true
|
|
|
|
|
|
si, ok := traffic.Strategies.Load(req.StrategyID)
|
|
|
if !ok {
|
|
|
// 不存在策略，就没有限制
|
|
|
return nil
|
|
|
}
|
|
|
s := si.(*misc.TrafficStrategy)
|
|
|
|
|
|
ari, ok := apiRates.Load(req.ApiName)
|
|
|
now := time.Now()
|
|
|
if !ok {
|
|
|
// 之前没有数据,初始化数据
|
|
|
ar := apiRate{1, now, 1, now, &sync.Map{}}
|
|
|
apiRates.Store(req.ApiName, &ar)
|
|
|
return nil
|
|
|
}
|
|
|
ar := ari.(*apiRate)
|
|
|
|
|
|
// QPS限制
|
|
|
if s.QPS != misc.STRATEGY_NO_LIMIT {
|
|
|
//要检查一下时间，如果超过一秒，则重新初始化数据
|
|
|
if now.Sub(ar.qpsUpdate) > 1e9 {
|
|
|
ar.qps = 1
|
|
|
ar.qpsUpdate = now
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
// 没超过一秒,超出限制，返回失败
|
|
|
if ar.qps >= s.QPS {
|
|
|
reply.Suc = false
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
// 没超过一秒，每超出，更新数据
|
|
|
ar.qps++
|
|
|
}
|
|
|
|
|
|
// 对并发进行限制
|
|
|
if s.Concurrent != misc.STRATEGY_NO_LIMIT {
|
|
|
// 这里要做一个防BUG设置，防止某些情况下，并发数到了最大值，结果不能减少或者归零
|
|
|
if ar.concurrent >= s.Concurrent {
|
|
|
// 如果当前时间距离上次更新并发时间，大于5秒，则认为出现了bug，重置并发
|
|
|
if now.Sub(ar.conUpdate) > 1e9 {
|
|
|
ar.concurrent = 0
|
|
|
ar.conUpdate = now
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
// 当前数量已经超过了并发数
|
|
|
reply.Suc = false
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
ar.concurrent++
|
|
|
ar.conUpdate = now
|
|
|
}
|
|
|
|
|
|
// 用户流量，访问次数限制
|
|
|
if req.ParamVal != "" {
|
|
|
pi, ok := ar.params.Load(req.ParamVal)
|
|
|
if !ok {
|
|
|
// 不存在记录，初始化
|
|
|
ar.params.Store(req.ParamVal, ¶m{1, now})
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
// 存在记录
|
|
|
p := pi.(*param)
|
|
|
// 判断当前时间和上次初始化的时间的跨度
|
|
|
if now.Sub(p.lastTime) > time.Duration(s.Span)*1e9 {
|
|
|
// 时间跨度已经超出，重新初始化
|
|
|
p.times = 1
|
|
|
p.lastTime = now
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
// 还在限制时间范围内,判断请求次数是否超出
|
|
|
if p.times >= s.Times {
|
|
|
reply.Suc = false
|
|
|
return nil
|
|
|
}
|
|
|
p.times++
|
|
|
}
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
func (rl *RateLimiter) DecApiRate(req misc.TrafficConReq, reply *misc.TrafficConRes) error {
|
|
|
reply.Suc = true
|
|
|
// 获取该api之前的计数
|
|
|
ari, ok := apiRates.Load(req.ApiName)
|
|
|
if !ok {
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
ar := ari.(*apiRate)
|
|
|
if ar.concurrent == 0 {
|
|
|
return nil
|
|
|
}
|
|
|
if ar.concurrent > 0 && ar.concurrent-1 <= 0 {
|
|
|
ar.concurrent = 0
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
ar.concurrent--
|
|
|
|
|
|
return nil
|
|
|
}
|
