新增方法解决并发添加redisconn问题

zhuCheer · Jul 24, 2020 · 1d6d861 · 1d6d861
1 parent 035e3bf
commit 1d6d861
Showing 1 changed file with 9 additions and 0 deletions.
diff --git 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 {