99 lines
2.1 KiB
Go
99 lines
2.1 KiB
Go
// Package rate_limit ...
|
|
//
|
|
// Description : rate_limit ...
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 2023-03-09 11:31
|
|
package rate_limit
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
redisInstance "github.com/go-redis/redis/v8"
|
|
"github.com/go-redis/redis_rate/v9"
|
|
)
|
|
|
|
var (
|
|
// limiter 限流实例
|
|
limiter *redis_rate.Limiter
|
|
)
|
|
|
|
// InitLimiter 初始化限流器实例
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 11:40 2023/3/9
|
|
func InitLimiter(redisClient *redisInstance.Client) {
|
|
limiter = redis_rate.NewLimiter(redisClient)
|
|
}
|
|
|
|
// Second 每秒允许的访问数
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 11:42 2023/3/9
|
|
func Second(ctx context.Context, key string, total, rate int) (bool, error) {
|
|
if total <= 0 || rate <= 0 {
|
|
return true, nil
|
|
}
|
|
res, err := limiter.AllowN(ctx, key, redis_rate.PerSecond(total), rate)
|
|
if nil != err {
|
|
return false, err
|
|
}
|
|
return res.Allowed > 0, nil
|
|
}
|
|
|
|
// Minute 每分钟允许访问数
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 11:46 2023/3/9
|
|
func Minute(ctx context.Context, key string, total, rate int) (bool, error) {
|
|
if total <= 0 || rate <= 0 {
|
|
return true, nil
|
|
}
|
|
res, err := limiter.AllowN(ctx, key, redis_rate.PerMinute(total), rate)
|
|
if nil != err {
|
|
return false, err
|
|
}
|
|
return res.Allowed > 0, nil
|
|
}
|
|
|
|
// Hour 每小时允许访问数
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 11:47 2023/3/9
|
|
func Hour(ctx context.Context, key string, total, rate int) (bool, error) {
|
|
if total <= 0 || rate <= 0 {
|
|
return true, nil
|
|
}
|
|
res, err := limiter.AllowN(ctx, key, redis_rate.PerHour(total), rate)
|
|
if nil != err {
|
|
return false, err
|
|
}
|
|
return res.Allowed > 0, nil
|
|
}
|
|
|
|
// Day 每天允许访问数
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 11:47 2023/3/9
|
|
func Day(ctx context.Context, key string, total, rate int) (bool, error) {
|
|
if total <= 0 || rate <= 0 {
|
|
return true, nil
|
|
}
|
|
res, err := 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
|
|
}
|