You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
go-zero/core/limit/periodlimit.go

127 lines
2.9 KiB
Go

4 years ago
package limit
import (
"context"
4 years ago
"errors"
"strconv"
"time"
"github.com/zeromicro/go-zero/core/stores/redis"
4 years ago
)
// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
const periodScript = `local limit = tonumber(ARGV[1])
4 years ago
local window = tonumber(ARGV[2])
local current = redis.call("INCRBY", KEYS[1], 1)
if current == 1 then
redis.call("expire", KEYS[1], window)
end
if current < limit then
4 years ago
return 1
elseif current == limit then
return 2
else
return 0
end`
const (
// Unknown means not initialized state.
4 years ago
Unknown = iota
// Allowed means allowed state.
4 years ago
Allowed
// HitQuota means this request exactly hit the quota.
4 years ago
HitQuota
// OverQuota means passed the quota.
4 years ago
OverQuota
internalOverQuota = 0
internalAllowed = 1
internalHitQuota = 2
)
// ErrUnknownCode is an error that represents unknown status code.
4 years ago
var ErrUnknownCode = errors.New("unknown status code")
type (
// PeriodOption defines the method to customize a PeriodLimit.
PeriodOption func(l *PeriodLimit)
4 years ago
// A PeriodLimit is used to limit requests during a period of time.
4 years ago
PeriodLimit struct {
period int
quota int
limitStore *redis.Redis
keyPrefix string
align bool
}
)
// NewPeriodLimit returns a PeriodLimit with given parameters.
4 years ago
func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string,
opts ...PeriodOption) *PeriodLimit {
4 years ago
limiter := &PeriodLimit{
period: period,
quota: quota,
limitStore: limitStore,
keyPrefix: keyPrefix,
}
for _, opt := range opts {
opt(limiter)
}
return limiter
}
// Take requests a permit, it returns the permit state.
4 years ago
func (h *PeriodLimit) Take(key string) (int, error) {
return h.TakeWithContext(context.Background(), key)
}
// TakeWithContext requests a permit with context, it returns the permit state.
func (h *PeriodLimit) TakeWithContext(ctx context.Context, key string) (int, error) {
resp, err := h.limitStore.EvalCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
4 years ago
strconv.Itoa(h.quota),
strconv.Itoa(h.calcExpireSeconds()),
})
if err != nil {
return Unknown, err
}
code, ok := resp.(int64)
if !ok {
return Unknown, ErrUnknownCode
}
switch code {
case internalOverQuota:
return OverQuota, nil
case internalAllowed:
return Allowed, nil
case internalHitQuota:
return HitQuota, nil
default:
return Unknown, ErrUnknownCode
}
}
func (h *PeriodLimit) calcExpireSeconds() int {
if h.align {
now := time.Now()
_, offset := now.Zone()
unix := now.Unix() + int64(offset)
4 years ago
return h.period - int(unix%int64(h.period))
}
return h.period
4 years ago
}
// Align returns a func to customize a PeriodLimit with alignment.
// For example, if we want to limit end users with 5 sms verification messages every day,
// we need to align with the local timezone and the start of the day.
func Align() PeriodOption {
4 years ago
return func(l *PeriodLimit) {
l.align = true
}
}