支持基于redis的流量控制 #1

Merged
zhangdeman merged 2 commits from feature/rate_limit into master 2023-03-09 14:06:43 +08:00
4 changed files with 102 additions and 0 deletions

1
.gitignore vendored
View File

@ -19,3 +19,4 @@
.idea
.vscode
mail_test.go
*_test.go

1
go.mod
View File

@ -26,6 +26,7 @@ require (
github.com/go-playground/locales v0.14.1 // indirect
github.com/go-playground/universal-translator v0.18.1 // indirect
github.com/go-playground/validator/v10 v10.11.2 // indirect
github.com/go-redis/redis_rate/v9 v9.1.2 // indirect
github.com/goccy/go-json v0.10.0 // indirect
github.com/json-iterator/go v1.1.12 // indirect
github.com/leodido/go-urn v1.2.1 // indirect

2
go.sum
View File

@ -66,6 +66,8 @@ github.com/go-playground/validator/v10 v10.11.2 h1:q3SHpufmypg+erIExEKUmsgmhDTyh
github.com/go-playground/validator/v10 v10.11.2/go.mod h1:NieE624vt4SCTJtD87arVLvdmjPAeV8BQlHtMnw9D7s=
github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI=
github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo=
github.com/go-redis/redis_rate/v9 v9.1.2 h1:H0l5VzoAtOE6ydd38j8MCq3ABlGLnvvbA1xDSVVCHgQ=
github.com/go-redis/redis_rate/v9 v9.1.2/go.mod h1:oam2de2apSgRG8aJzwJddXbNu91Iyz1m8IKJE2vpvlQ=
github.com/go-task/slim-sprig v0.0.0-20210107165309-348f09dbbbc0/go.mod h1:fyg7847qk6SyHyPtNmDHnmrv/HOrqktSC+C9fM+CJOE=
github.com/goccy/go-json v0.9.7 h1:IcB+Aqpx/iMHu5Yooh7jEzJk1JZ7Pjtmys2ukPr7EeM=
github.com/goccy/go-json v0.9.7/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=

98
rate_limit/limit.go Normal file
View File

@ -0,0 +1,98 @@
// 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
}