// 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/define" "github.com/go-redis/redis_rate/v10" redisInstance "github.com/redis/go-redis/v9" "sync" "time" ) var ( RedisClient = &Redis{} ) // Redis ... // // Author : go_developer@163.com<白茶清欢> // // Date : 11:44 2024/6/20 type Redis struct { limiter *redis_rate.Limiter // limiter 限流实例 redisClient *redisInstance.ClusterClient // redis 客户端实例 lock *sync.RWMutex } 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 } 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 } */