zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

chore: refactor periodlimit (#1428)

Browse Source 
* chore: refactor periodlimit

* chore: add comments
This commit is contained in:
Kevin Wan
2022-01-09 16:22:34 +08:00
committed by GitHub
parent 5e0e2d2b14
commit 53af194ef9
1 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
core/limit/periodlimit.go
View File

@@ -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