Эх сурвалжийг харах

periodlimit new function TakeWithContext (#1983)

Co-authored-by: mashouyue's m1max <mashouyue@toowow.cn>
马守越 2 жил өмнө
parent
commit
392a390a3f

+ 7 - 1
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()),
 	})