|
|
@ -8,9 +8,8 @@ import (
|
|
|
|
"github.com/zeromicro/go-zero/core/stores/redis"
|
|
|
|
"github.com/zeromicro/go-zero/core/stores/redis"
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
const (
|
|
|
|
// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
|
|
|
|
// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key
|
|
|
|
const periodScript = `local limit = tonumber(ARGV[1])
|
|
|
|
periodScript = `local limit = tonumber(ARGV[1])
|
|
|
|
|
|
|
|
local window = tonumber(ARGV[2])
|
|
|
|
local window = tonumber(ARGV[2])
|
|
|
|
local current = redis.call("INCRBY", KEYS[1], 1)
|
|
|
|
local current = redis.call("INCRBY", KEYS[1], 1)
|
|
|
|
if current == 1 then
|
|
|
|
if current == 1 then
|
|
|
@ -23,8 +22,6 @@ elseif current == limit then
|
|
|
|
else
|
|
|
|
else
|
|
|
|
return 0
|
|
|
|
return 0
|
|
|
|
end`
|
|
|
|
end`
|
|
|
|
zoneDiff = 3600 * 8 // GMT+8 for our services
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
const (
|
|
|
|
const (
|
|
|
|
// Unknown means not initialized state.
|
|
|
|
// Unknown means not initialized state.
|
|
|
@ -104,7 +101,9 @@ func (h *PeriodLimit) Take(key string) (int, error) {
|
|
|
|
|
|
|
|
|
|
|
|
func (h *PeriodLimit) calcExpireSeconds() int {
|
|
|
|
func (h *PeriodLimit) calcExpireSeconds() int {
|
|
|
|
if h.align {
|
|
|
|
if h.align {
|
|
|
|
unix := time.Now().Unix() + zoneDiff
|
|
|
|
now := time.Now()
|
|
|
|
|
|
|
|
_, offset := now.Zone()
|
|
|
|
|
|
|
|
unix := now.Unix() + int64(offset)
|
|
|
|
return h.period - int(unix%int64(h.period))
|
|
|
|
return h.period - int(unix%int64(h.period))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
@ -112,6 +111,8 @@ func (h *PeriodLimit) calcExpireSeconds() int {
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// Align returns a func to customize a PeriodLimit with alignment.
|
|
|
|
// 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 {
|
|
|
|
func Align() PeriodOption {
|
|
|
|
return func(l *PeriodLimit) {
|
|
|
|
return func(l *PeriodLimit) {
|
|
|
|
l.align = true
|
|
|
|
l.align = true
|
|
|
|