|
@@ -27,9 +27,13 @@ end`
|
|
|
)
|
|
|
|
|
|
const (
|
|
|
+ // Unknown means not initialized state.
|
|
|
Unknown = iota
|
|
|
+ // Allowed means allowed state.
|
|
|
Allowed
|
|
|
+ // HitQuota means this request exactly hit the quota.
|
|
|
HitQuota
|
|
|
+ // OverQuota means passed the quota.
|
|
|
OverQuota
|
|
|
|
|
|
internalOverQuota = 0
|
|
@@ -37,11 +41,14 @@ const (
|
|
|
internalHitQuota = 2
|
|
|
)
|
|
|
|
|
|
+// ErrUnknownCode is an error that represents unknown status code.
|
|
|
var ErrUnknownCode = errors.New("unknown status code")
|
|
|
|
|
|
type (
|
|
|
- LimitOption func(l *PeriodLimit)
|
|
|
+ // PeriodOption defines the method to customize a PeriodLimit.
|
|
|
+ PeriodOption func(l *PeriodLimit)
|
|
|
|
|
|
+ // A PeriodLimit is used to limit requests during a period of time.
|
|
|
PeriodLimit struct {
|
|
|
period int
|
|
|
quota int
|
|
@@ -51,8 +58,9 @@ type (
|
|
|
}
|
|
|
)
|
|
|
|
|
|
+// NewPeriodLimit returns a PeriodLimit with given parameters.
|
|
|
func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string,
|
|
|
- opts ...LimitOption) *PeriodLimit {
|
|
|
+ opts ...PeriodOption) *PeriodLimit {
|
|
|
limiter := &PeriodLimit{
|
|
|
period: period,
|
|
|
quota: quota,
|
|
@@ -67,6 +75,7 @@ func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string
|
|
|
return limiter
|
|
|
}
|
|
|
|
|
|
+// Take requests a permit, it returns the permit state.
|
|
|
func (h *PeriodLimit) Take(key string) (int, error) {
|
|
|
resp, err := h.limitStore.Eval(periodScript, []string{h.keyPrefix + key}, []string{
|
|
|
strconv.Itoa(h.quota),
|
|
@@ -102,7 +111,8 @@ func (h *PeriodLimit) calcExpireSeconds() int {
|
|
|
return h.period
|
|
|
}
|
|
|
|
|
|
-func Align() LimitOption {
|
|
|
+// Align returns a func to customize a PeriodLimit with alignment.
|
|
|
+func Align() PeriodOption {
|
|
|
return func(l *PeriodLimit) {
|
|
|
l.align = true
|
|
|
}
|