token limit support context (#2335)

* token limit support context

* add token limit with ctx

add token limit with ctx

Co-authored-by: sado <liaoyonglin@bilibili.com>

core/limit/tokenlimit_test.go

@@ -1,6 +1,7 @@
 package limit
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -15,6 +16,30 @@ func init() {
 	logx.Disable()
 }
 
+func TestTokenLimit_WithCtx(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+
+	const (
+		total = 100
+		rate  = 5
+		burst = 10
+	)
+	l := NewTokenLimiter(rate, burst, redis.New(s.Addr()), "tokenlimit")
+	defer s.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	ok := l.AllowCtx(ctx)
+	assert.True(t, ok)
+
+	cancel()
+	for i := 0; i < total; i++ {
+		ok := l.AllowCtx(ctx)
+		assert.False(t, ok)
+		assert.False(t, l.monitorStarted)
+	}
+}
+
 func TestTokenLimit_Rescue(t *testing.T) {
 	s, err := miniredis.Run()
 	assert.Nil(t, err)