流控接口定义 + 基于redis的流控实现

redis.go

@@ -0,0 +1,130 @@
+// Package rate_limit ...
+//
+// Description : rate_limit ...
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 2023-03-09 11:31
+package rate_limit
+
+import (
+	"context"
+	"git.zhangdeman.cn/zhangdeman/rate_limit/abstract"
+	"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),
+	}
+}
+
+// Redis ...
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 11:44 2024/6/20
+type Redis struct {
+	limiter *redis_rate.Limiter // limiter 限流实例
+}
+
+func (r *Redis) Second(ctx context.Context, key string, total, rate int) (bool, error) {
+	if total <= 0 || rate <= 0 {
+		return true, nil
+	}
+	res, err := r.limiter.AllowN(ctx, key, redis_rate.PerSecond(total), rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}
+
+func (r *Redis) Minute(ctx context.Context, key string, total, rate int) (bool, error) {
+	if total <= 0 || rate <= 0 {
+		return true, nil
+	}
+	res, err := r.limiter.AllowN(ctx, key, redis_rate.PerMinute(total), rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}
+
+func (r *Redis) Hour(ctx context.Context, key string, total, rate int) (bool, error) {
+	if total <= 0 || rate <= 0 {
+		return true, nil
+	}
+	res, err := r.limiter.AllowN(ctx, key, redis_rate.PerHour(total), rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}
+
+func (r *Redis) Day(ctx context.Context, key string, total, rate int) (bool, error) {
+	if total <= 0 || rate <= 0 {
+		return true, nil
+	}
+	res, err := r.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
+}
+
+func (r *Redis) Month(ctx context.Context, key string, total, rate int) (bool, error) {
+	if total <= 0 || rate <= 0 {
+		return true, nil
+	}
+	res, err := r.limiter.AllowN(ctx, key, redis_rate.Limit{
+		Rate:   rate,
+		Period: 24 * 30 * time.Hour,
+		Burst:  total,
+	}, rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}
+
+func (r *Redis) Year(ctx context.Context, key string, total, rate int) (bool, error) {
+	if total <= 0 || rate <= 0 {
+		return true, nil
+	}
+	res, err := r.limiter.AllowN(ctx, key, redis_rate.Limit{
+		Rate:   rate,
+		Period: 24 * 30 * 365 * time.Hour,
+		Burst:  total,
+	}, rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}
+
+func (r *Redis) Custom(ctx context.Context, timeInfo int64, key string, total, rate int) (bool, error) {
+	if total <= 0 || rate <= 0 {
+		return true, nil
+	}
+	res, err := r.limiter.AllowN(ctx, key, redis_rate.Limit{
+		Rate:   rate,
+		Period: time.Duration(timeInfo) * time.Second,
+		Burst:  total,
+	}, rate)
+	if nil != err {
+		return false, err
+	}
+	return res.Allowed > 0, nil
+}