This commit is contained in:
jiangyong 2026-05-08 18:33:50 +08:00
parent 3559f07dfe
commit d1433040df
1 changed files with 14 additions and 14 deletions

View File

@ -9,9 +9,9 @@ import (
) )
type WindowsLimitOption struct { type WindowsLimitOption struct {
keyPrefix string //redis可以前缀会拼接上succ和fail Key string //redis可以前缀会拼接上succ和fail
limitCount int64 //限制数量 Count int64 //限制数量
scope time.Duration //限制时间区间,默认一分钟 Size time.Duration //限制时间区间,默认一分钟
succKey string succKey string
failKey string failKey string
} }
@ -22,9 +22,9 @@ type WindowsLimit struct {
func NewWindowsLimitOption() *WindowsLimitOption { func NewWindowsLimitOption() *WindowsLimitOption {
return &WindowsLimitOption{ return &WindowsLimitOption{
keyPrefix: fmt.Sprintf("windows:limit:%s:%d", time.Now().Format("20060102"), time.Now().UnixMilli()), Key: fmt.Sprintf("windows:limit:%s:%d", time.Now().Format("20060102"), time.Now().UnixMilli()),
limitCount: 1000, Count: 1000,
scope: time.Second * 60, Size: time.Second * 60,
} }
} }
@ -37,7 +37,7 @@ func NewWindowsLimit(client *redis.Client, option *WindowsLimitOption) *WindowsL
func (w *WindowsLimit) tidy() { func (w *WindowsLimit) tidy() {
now := time.Now().UnixMilli() now := time.Now().UnixMilli()
start := now - w.option.scope.Milliseconds() start := now - w.option.Size.Milliseconds()
pipe := w.client.Pipeline() pipe := w.client.Pipeline()
pipe.ZRemRangeByScore(w.option.failKey, "0", fmt.Sprintf("%d", start)) pipe.ZRemRangeByScore(w.option.failKey, "0", fmt.Sprintf("%d", start))
@ -47,7 +47,7 @@ func (w *WindowsLimit) tidy() {
func (w *WindowsLimit) Count() (int64, int64, error) { func (w *WindowsLimit) Count() (int64, int64, error) {
now := time.Now().UnixMilli() now := time.Now().UnixMilli()
start := now - w.option.scope.Milliseconds() start := now - w.option.Size.Milliseconds()
pipe := w.client.Pipeline() pipe := w.client.Pipeline()
succCmd := pipe.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now)) succCmd := pipe.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
@ -61,7 +61,7 @@ func (w *WindowsLimit) Count() (int64, int64, error) {
func (w *WindowsLimit) Allow() bool { func (w *WindowsLimit) Allow() bool {
now := time.Now().UnixMilli() now := time.Now().UnixMilli()
start := now - w.option.scope.Milliseconds() start := now - w.option.Size.Milliseconds()
pipe := w.client.Pipeline() pipe := w.client.Pipeline()
succCmd := w.client.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now)) succCmd := w.client.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
_, err := pipe.Exec() _, err := pipe.Exec()
@ -69,7 +69,7 @@ func (w *WindowsLimit) Allow() bool {
return false return false
} }
if succCmd.Val() >= w.option.limitCount { if succCmd.Val() >= w.option.Count {
return false return false
} }
return true return true
@ -85,8 +85,8 @@ func (w *WindowsLimit) Add(succ bool) error {
pipe.ZAdd(w.option.failKey, redis.Z{Score: float64(now), Member: fmt.Sprintf("%d:%d", now, rand.Int())}) pipe.ZAdd(w.option.failKey, redis.Z{Score: float64(now), Member: fmt.Sprintf("%d:%d", now, rand.Int())})
} }
pipe.Expire(w.option.failKey, w.option.scope) pipe.Expire(w.option.failKey, w.option.Size)
pipe.Expire(w.option.succKey, w.option.scope) pipe.Expire(w.option.succKey, w.option.Size)
_, err := pipe.Exec() _, err := pipe.Exec()
return err return err
} }