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 }