forked from zeromicro/go-zero
/
periodlimit.go
109 lines (93 loc) · 1.97 KB
/
periodlimit.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
package limit
import (
"errors"
"strconv"
"time"
"github.com/dylanNew/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])
local window = tonumber(ARGV[2])
local current = redis.call("INCRBY", KEYS[1], 1)
if current == 1 then
redis.call("expire", KEYS[1], window)
return 1
elseif current < limit then
return 1
elseif current == limit then
return 2
else
return 0
end`
zoneDiff = 3600 * 8 // GMT+8 for our services
)
const (
Unknown = iota
Allowed
HitQuota
OverQuota
internalOverQuota = 0
internalAllowed = 1
internalHitQuota = 2
)
var ErrUnknownCode = errors.New("unknown status code")
type (
LimitOption func(l *PeriodLimit)
PeriodLimit struct {
period int
quota int
limitStore *redis.Redis
keyPrefix string
align bool
}
)
func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string,
opts ...LimitOption) *PeriodLimit {
limiter := &PeriodLimit{
period: period,
quota: quota,
limitStore: limitStore,
keyPrefix: keyPrefix,
}
for _, opt := range opts {
opt(limiter)
}
return limiter
}
func (h *PeriodLimit) Take(key string) (int, error) {
resp, err := h.limitStore.Eval(periodScript, []string{h.keyPrefix + key}, []string{
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 {
unix := time.Now().Unix() + zoneDiff
return h.period - int(unix%int64(h.period))
} else {
return h.period
}
}
func Align() LimitOption {
return func(l *PeriodLimit) {
l.align = true
}
}