feat: code cleanup
This commit is contained in:
@@ -9,14 +9,11 @@ package abstract
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
|
||||||
"git.zhangdeman.cn/zhangdeman/rate_limit/define"
|
"git.zhangdeman.cn/zhangdeman/rate_limit/define"
|
||||||
)
|
)
|
||||||
|
|
||||||
// IRateLimit 流控实现接口约束
|
// IRateLimit 流控实现接口约束
|
||||||
//
|
|
||||||
// Author : go_developer@163.com<白茶清欢>
|
|
||||||
//
|
|
||||||
// Date : 11:33 2024/6/20
|
|
||||||
type IRateLimit interface {
|
type IRateLimit interface {
|
||||||
AllowN(ctx context.Context, limitCfg *define.LimitConfig, tokenCnt int) (bool, error) // 申请N个令牌, N >= 0
|
AllowN(ctx context.Context, limitCfg *define.LimitConfig, tokenCnt int) (bool, error) // 申请N个令牌, N >= 0
|
||||||
Reset(ctx context.Context, limitCfg *define.LimitConfig) error // 重置
|
Reset(ctx context.Context, limitCfg *define.LimitConfig) error // 重置
|
||||||
|
|||||||
@@ -8,10 +8,6 @@
|
|||||||
package define
|
package define
|
||||||
|
|
||||||
// LimitConfig ...
|
// LimitConfig ...
|
||||||
//
|
|
||||||
// Author : go_developer@163.com<白茶清欢>
|
|
||||||
//
|
|
||||||
// Date : 15:17 2024/6/20
|
|
||||||
type LimitConfig struct {
|
type LimitConfig struct {
|
||||||
Key string `json:"key"` // 操作的key
|
Key string `json:"key"` // 操作的key
|
||||||
Total int `json:"total"` // 令牌桶大小
|
Total int `json:"total"` // 令牌桶大小
|
||||||
|
|||||||
@@ -9,10 +9,11 @@ package rate_limit
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"git.zhangdeman.cn/zhangdeman/rate_limit/define"
|
|
||||||
"golang.org/x/time/rate"
|
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"git.zhangdeman.cn/zhangdeman/rate_limit/define"
|
||||||
|
"golang.org/x/time/rate"
|
||||||
)
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
@@ -23,10 +24,6 @@ var (
|
|||||||
)
|
)
|
||||||
|
|
||||||
// Memory ...
|
// Memory ...
|
||||||
//
|
|
||||||
// Author : go_developer@163.com<白茶清欢>
|
|
||||||
//
|
|
||||||
// Date : 20:40 2024/6/20
|
|
||||||
type Memory struct {
|
type Memory struct {
|
||||||
lock *sync.RWMutex
|
lock *sync.RWMutex
|
||||||
limiterTable map[string]*rate.Limiter
|
limiterTable map[string]*rate.Limiter
|
||||||
|
|||||||
99
redis.go
99
redis.go
@@ -9,11 +9,12 @@ package rate_limit
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
"git.zhangdeman.cn/zhangdeman/rate_limit/define"
|
"git.zhangdeman.cn/zhangdeman/rate_limit/define"
|
||||||
"github.com/go-redis/redis_rate/v10"
|
"github.com/go-redis/redis_rate/v10"
|
||||||
redisInstance "github.com/redis/go-redis/v9"
|
redisInstance "github.com/redis/go-redis/v9"
|
||||||
"sync"
|
|
||||||
"time"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
@@ -54,97 +55,3 @@ func (r *Redis) AllowN(ctx context.Context, limitCfg *define.LimitConfig, tokenC
|
|||||||
func (r *Redis) Reset(ctx context.Context, limitCfg *define.LimitConfig) error {
|
func (r *Redis) Reset(ctx context.Context, limitCfg *define.LimitConfig) error {
|
||||||
return r.limiter.Reset(ctx, limitCfg.Key)
|
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
|
|
||||||
}
|
|
||||||
*/
|
|
||||||
|
|||||||
Reference in New Issue
Block a user