github_mirrors/go-zero
1
0
Fork 0
mirror of https://github.com/zeromicro/go-zero.git synced 2025-01-23 09:00:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

periodlimit new function TakeWithContext (#1983)

Browse Source 
Co-authored-by: mashouyue's m1max <mashouyue@toowow.cn>
This commit is contained in:
马守越 2022-06-11 12:07:57 +08:00 committed by GitHub
parent 2a900e1795
commit 392a390a3f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
core/limit/periodlimit.go
View File

@ -1,6 +1,7 @@
package limit
import (
"context"
"errors"
"strconv"
"time"
@ -74,7 +75,12 @@ func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string
// 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{
return h.TakeWithContext(context.Background(), key)
}
// TakeWithContext requests a permit with context, it returns the permit state.
func (h *PeriodLimit) TakeWithContext(ctx context.Context, key string) (int, error) {
resp, err := h.limitStore.EvalCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
strconv.Itoa(h.quota),
strconv.Itoa(h.calcExpireSeconds()),
})