rate_limit/redis.go

// 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.Client // redis 客户端实例
	lock        *sync.RWMutex
}

func (r *Redis) SetRedisClient(client *redisInstance.Client) {
	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
}
*/