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
|
|||
|
|
}
|