chore: refactor periodlimit (#1428)

* chore: refactor periodlimit

* chore: add comments
master
Kevin Wan 3 years ago committed by GitHub
parent 5e0e2d2b14
commit 53af194ef9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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
periodScript = `local limit = tonumber(ARGV[1]) const 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

Loading…
Cancel
Save