1

pkg/utility/ratecheck/ratecheck.go

@@ -0,0 +1,47 @@
+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
+}