From 53af194ef9fff9d22ef8e4e0161d0ce36e05cdad Mon Sep 17 00:00:00 2001 From: Kevin Wan Date: Sun, 9 Jan 2022 16:22:34 +0800 Subject: [PATCH] chore: refactor periodlimit (#1428) * chore: refactor periodlimit * chore: add comments --- core/limit/periodlimit.go | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/core/limit/periodlimit.go b/core/limit/periodlimit.go index 1dca26c0..edb737e8 100644 --- a/core/limit/periodlimit.go +++ b/core/limit/periodlimit.go @@ -8,9 +8,8 @@ import ( "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 - periodScript = `local limit = tonumber(ARGV[1]) +// to be compatible with aliyun redis, we cannot use `local key = KEYS[1]` to reuse the key +const periodScript = `local limit = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local current = redis.call("INCRBY", KEYS[1], 1) if current == 1 then @@ -23,8 +22,6 @@ elseif current == limit then else return 0 end` - zoneDiff = 3600 * 8 // GMT+8 for our services -) const ( // Unknown means not initialized state. @@ -104,7 +101,9 @@ func (h *PeriodLimit) Take(key string) (int, error) { func (h *PeriodLimit) calcExpireSeconds() int { 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)) } @@ -112,6 +111,8 @@ func (h *PeriodLimit) calcExpireSeconds() int { } // 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 { return func(l *PeriodLimit) { l.align = true