97 lines
2.4 KiB
Go
97 lines
2.4 KiB
Go
package goutil
|
||
|
||
import (
|
||
"fmt"
|
||
"math/rand"
|
||
"time"
|
||
|
||
"github.com/go-redis/redis"
|
||
)
|
||
|
||
type WindowsLimitOption struct {
|
||
Key string //redis可以前缀，会拼接上succ和fail
|
||
Count int64 //限制数量
|
||
Size time.Duration //限制时间区间，默认一分钟
|
||
succKey string
|
||
failKey string
|
||
}
|
||
type WindowsLimit struct {
|
||
client *redis.Client
|
||
option *WindowsLimitOption
|
||
}
|
||
|
||
func NewWindowsLimit(client *redis.Client, option *WindowsLimitOption) *WindowsLimit {
|
||
option.Default()
|
||
return &WindowsLimit{
|
||
client: client,
|
||
option: option,
|
||
}
|
||
}
|
||
|
||
func (o *WindowsLimitOption) Default() {
|
||
if o.Key == "" {
|
||
o.Key = fmt.Sprintf("windows:limit:%s:%d", time.Now().Format("20060102"), time.Now().UnixMilli())
|
||
}
|
||
if o.Count <= 0 {
|
||
o.Count = 1000
|
||
}
|
||
if o.Size <= 0 {
|
||
o.Size = time.Second * 60
|
||
}
|
||
}
|
||
func (w *WindowsLimit) tidy() {
|
||
now := time.Now().UnixMilli()
|
||
start := now - w.option.Size.Milliseconds()
|
||
|
||
pipe := w.client.Pipeline()
|
||
pipe.ZRemRangeByScore(w.option.failKey, "0", fmt.Sprintf("%d", start))
|
||
pipe.ZRemRangeByScore(w.option.succKey, "0", fmt.Sprintf("%d", start))
|
||
pipe.Exec()
|
||
}
|
||
|
||
func (w *WindowsLimit) Count() (int64, int64, error) {
|
||
now := time.Now().UnixMilli()
|
||
start := now - w.option.Size.Milliseconds()
|
||
|
||
pipe := w.client.Pipeline()
|
||
succCmd := pipe.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
|
||
failCmd := pipe.ZCount(w.option.failKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
|
||
_, err := pipe.Exec()
|
||
if err != nil {
|
||
return 0, 0, err
|
||
}
|
||
return succCmd.Val(), failCmd.Val(), nil
|
||
}
|
||
|
||
func (w *WindowsLimit) Allow() bool {
|
||
now := time.Now().UnixMilli()
|
||
start := now - w.option.Size.Milliseconds()
|
||
pipe := w.client.Pipeline()
|
||
succCmd := w.client.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
|
||
_, err := pipe.Exec()
|
||
if err != nil {
|
||
return true
|
||
}
|
||
|
||
if succCmd.Val() < w.option.Count {
|
||
return true
|
||
}
|
||
return false
|
||
}
|
||
|
||
func (w *WindowsLimit) Add(succ bool) error {
|
||
now := time.Now().UnixMilli()
|
||
w.tidy()
|
||
pipe := w.client.Pipeline()
|
||
if succ {
|
||
pipe.ZAdd(w.option.succKey, redis.Z{Score: float64(now), Member: fmt.Sprintf("%d:%d", now, rand.Int())})
|
||
} else {
|
||
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.Size)
|
||
pipe.Expire(w.option.succKey, w.option.Size)
|
||
_, err := pipe.Exec()
|
||
return err
|
||
}
|