From a43edd511b6d37c82b3fd3f4d5752fe8ac7a1a3a Mon Sep 17 00:00:00 2001 From: mashouyue's m1max Date: Tue, 7 Jun 2022 11:43:03 +0800 Subject: [PATCH] periodlimit new function TakeWithContext --- 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 73fb9d009401..20ad93fb8b15 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()), })