redis流控 + 单元测试

redis.go

@@ -9,23 +9,16 @@ package rate_limit
 
 import (
 	"context"
-	"git.zhangdeman.cn/zhangdeman/rate_limit/abstract"
+	"git.zhangdeman.cn/zhangdeman/rate_limit/define"
+	"github.com/go-redis/redis_rate/v10"
+	redisInstance "github.com/redis/go-redis/v9"
+	"sync"
 	"time"
-
-	redisInstance "github.com/go-redis/redis/v8"
-	"github.com/go-redis/redis_rate/v9"
 )
 
-// NewRedis redis流控实例
-//
-// Author : go_developer@163.com<白茶清欢>
-//
-// Date : 11:43 2024/6/20
-func NewRedis(redisClient *redisInstance.Client) abstract.IRateLimit {
-	return &Redis{
-		limiter: redis_rate.NewLimiter(redisClient),
-	}
-}
+var (
+	RedisClient = &Redis{}
+)
 
 // Redis ...
 //
@@ -33,10 +26,36 @@ func NewRedis(redisClient *redisInstance.Client) abstract.IRateLimit {
 //
 // Date : 11:44 2024/6/20
 type Redis struct {
-	limiter *redis_rate.Limiter // limiter 限流实例
+	limiter     *redis_rate.Limiter          // limiter 限流实例
+	redisClient *redisInstance.ClusterClient // redis 客户端实例
+	lock        *sync.RWMutex
 }
 
-func (r *Redis) Second(ctx context.Context, key string, total, rate int) (bool, error) {
+func (r *Redis) SetRedisClient(client *redisInstance.ClusterClient) {
+	r.redisClient = client
+	r.limiter = redis_rate.NewLimiter(r.redisClient)
+}
+
+func (r *Redis) AllowN(ctx context.Context, limitCfg *define.LimitConfig, tokenCnt int) (bool, error) {
+	if limitCfg.Total <= 0 || limitCfg.Rate <= 0 {
+		return true, nil
+	}
+	res, err := r.limiter.AllowN(ctx, limitCfg.Key, redis_rate.Limit{
+		Rate:   limitCfg.Rate,
+		Period: time.Duration(limitCfg.TimeInterval) * time.Second,
+		Burst:  limitCfg.Total,
+	}, tokenCnt)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}
+
+func (r *Redis) Reset(ctx context.Context, limitCfg *define.LimitConfig) error {
+	return r.limiter.Reset(ctx, limitCfg.Key)
+}
+
+/*func (r *Redis) Second(ctx context.Context, key string, total, rate int) (bool, error) {
 	if total <= 0 || rate <= 0 {
 		return true, nil
 	}
@@ -128,3 +147,4 @@ func (r *Redis) Custom(ctx context.Context, timeInfo int64, key string, total, r
 	}
 	return res.Allowed > 0, nil
 }
+*/