48 lines
1.2 KiB
Go
48 lines
1.2 KiB
Go
package ratecheck
|
||
|
||
import (
|
||
"github.com/juju/ratelimit"
|
||
gocache "github.com/patrickmn/go-cache"
|
||
"go-admin/pkg/utility"
|
||
"time"
|
||
)
|
||
|
||
var (
|
||
// Map of limiters with TTL
|
||
tokenBuckets = gocache.New(120*time.Minute, 1*time.Minute)
|
||
userDur = 1 * time.Second
|
||
userSize int64 = 1
|
||
|
||
orderDur = 20 * time.Second
|
||
)
|
||
|
||
// CheckRateLimit 根据key检测，在规定时间内，是否超过访问次数,限流
|
||
func CheckRateLimit(key string, duration time.Duration, size int64) bool {
|
||
if _, found := tokenBuckets.Get(key); !found {
|
||
tokenBuckets.Set(
|
||
key,
|
||
ratelimit.NewBucket(duration, size),
|
||
duration)
|
||
}
|
||
expiringMap, found := tokenBuckets.Get(key)
|
||
if !found {
|
||
return false
|
||
}
|
||
return expiringMap.(*ratelimit.Bucket).TakeAvailable(1) > 0
|
||
}
|
||
|
||
// CheckUserRateLimit 根据key检测，在规定时间内，单个用户是否超过访问次数,限流,默认5秒1次请求
|
||
func CheckUserRateLimit(userid int, methodName string) bool {
|
||
key := methodName + "-" + utility.IntTostring(userid)
|
||
return CheckRateLimit(key, userDur, userSize)
|
||
}
|
||
|
||
// 检测订单成交是否重复推送
|
||
func CheckOrderIdIsExist(tradeId string) bool {
|
||
_, found := tokenBuckets.Get(tradeId)
|
||
if !found {
|
||
tokenBuckets.Set(tradeId, true, orderDur)
|
||
}
|
||
return found
|
||
}
|