支持秒、分钟、小时、天级别的流控

rate_limit/limit.go

@@ -0,0 +1,86 @@
+// Package rate_limit ...
+//
+// Description : rate_limit ...
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 2023-03-09 11:31
+package rate_limit
+
+import (
+	"context"
+	"time"
+
+	redisInstance "github.com/go-redis/redis/v8"
+	"github.com/go-redis/redis_rate/v9"
+)
+
+var (
+	// limiter 限流实例
+	limiter *redis_rate.Limiter
+)
+
+// InitLimiter 初始化限流器实例
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 11:40 2023/3/9
+func InitLimiter(redisClient *redisInstance.Client) {
+	limiter = redis_rate.NewLimiter(redisClient)
+}
+
+// Second 每秒允许的访问数
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 11:42 2023/3/9
+func Second(ctx context.Context, key string, total, rate int) bool {
+	res, err := limiter.AllowN(ctx, key, redis_rate.PerSecond(total), rate)
+	if nil != err {
+		return false
+	}
+	return res.Allowed > 0
+}
+
+// Minute 每分钟允许访问数
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 11:46 2023/3/9
+func Minute(ctx context.Context, key string, total, rate int) (bool, error) {
+	res, err := limiter.AllowN(ctx, key, redis_rate.PerMinute(total), rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}
+
+// Hour 每小时允许访问数
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 11:47 2023/3/9
+func Hour(ctx context.Context, key string, total, rate int) (bool, error) {
+	res, err := limiter.AllowN(ctx, key, redis_rate.PerHour(total), rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}
+
+// Day 每天允许访问数
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 11:47 2023/3/9
+func Day(ctx context.Context, key string, total, rate int) (bool, error) {
+	res, err := limiter.AllowN(ctx, key, redis_rate.Limit{
+		Rate:   rate,
+		Period: 24 * time.Hour,
+		Burst:  total,
+	}, rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}