151 lines
3.5 KiB
Go
151 lines
3.5 KiB
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.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
|
|
}
|
|
*/
|