From 392a390a3f9c8118d028f6331a56ca70bd579f32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=A9=AC=E5=AE=88=E8=B6=8A?= Date: Sat, 11 Jun 2022 12:07:57 +0800 Subject: [PATCH] periodlimit new function TakeWithContext (#1983) Co-authored-by: mashouyue's m1max --- core/limit/periodlimit.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/core/limit/periodlimit.go b/core/limit/periodlimit.go index 73fb9d00..20ad93fb 100644 --- a/core/limit/periodlimit.go +++ b/core/limit/periodlimit.go @@ -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()), })