From 1d6d861c562bbed392be7cbe78117774aac616e3 Mon Sep 17 00:00:00 2001
From: zhucheer <zhu.cheer@qq.com>
Date: Fri, 24 Jul 2020 23:32:19 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=E6=96=B9=E6=B3=95=E8=A7=A3?=
 =?UTF-8?q?=E5=86=B3=E5=B9=B6=E5=8F=91=E6=B7=BB=E5=8A=A0redisconn=E9=97=AE?=
 =?UTF-8?q?=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ratelimit.go | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/ratelimit.go b/ratelimit.go
index c57057d..20ab6ba 100644
--- a/ratelimit.go
+++ b/ratelimit.go
@@ -242,6 +242,15 @@ func (tb *Bucket) TakeAvailable(count int64) int64 {
 	return tb.takeAvailable(tb.clock.Now(), count)
 }
 
+func (tb *Bucket) TakeAvailableWaitWithRedis(count int64, maxWait time.Duration, conn redis.Conn) (token int64) {
+	tb.mu.Lock()
+	defer tb.mu.Unlock()
+	tb.redisConn = conn
+
+	return tb.TakeAvailableWait(count, maxWait)
+}
+
+
 func (tb *Bucket) TakeAvailableWait(count int64, maxWait time.Duration) (token int64) {
 	timer := time.NewTimer(maxWait)
 	for {