support not wait take func

limiter_atomic.go

@@ -21,9 +21,8 @@
 package ratelimit // import "go.uber.org/ratelimit"
 
 import (
-	"time"
-
 	"sync/atomic"
+	"time"
 	"unsafe"
 )
 
@@ -66,6 +65,16 @@ func newAtomicBased(rate int, opts ...Option) *atomicLimiter {
 // Take blocks to ensure that the time spent between multiple
 // Take calls is on average per/rate.
 func (t *atomicLimiter) Take() time.Time {
+	last, interval := t.take()
+	t.clock.Sleep(interval)
+	return last
+}
+
+func (t *atomicLimiter) TakeNoWait() (time.Time, time.Duration) {
+	return t.take()
+}
+
+func (t *atomicLimiter) take() (time.Time, time.Duration) {
 	var (
 		newState state
 		taken    bool
@@ -105,6 +114,5 @@ func (t *atomicLimiter) Take() time.Time {
 		}
 		taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
 	}
-	t.clock.Sleep(interval)
-	return newState.last
+	return newState.last, interval
 }

limiter_atomic_int64.go

@@ -21,9 +21,8 @@
 package ratelimit // import "go.uber.org/ratelimit"
 
 import (
-	"time"
-
 	"sync/atomic"
+	"time"
 )
 
 type atomicInt64Limiter struct {
@@ -57,6 +56,16 @@ func newAtomicInt64Based(rate int, opts ...Option) *atomicInt64Limiter {
 // Take blocks to ensure that the time spent between multiple
 // Take calls is on average time.Second/rate.
 func (t *atomicInt64Limiter) Take() time.Time {
+	last, interval := t.take()
+	t.clock.Sleep(interval)
+	return last
+}
+
+func (t *atomicInt64Limiter) TakeNoWait() (time.Time, time.Duration) {
+	return t.take()
+}
+
+func (t *atomicInt64Limiter) take() (time.Time, time.Duration) {
 	var (
 		newTimeOfNextPermissionIssue int64
 		now                          int64
@@ -82,6 +91,5 @@ func (t *atomicInt64Limiter) Take() time.Time {
 			break
 		}
 	}
-	t.clock.Sleep(time.Duration(newTimeOfNextPermissionIssue - now))
-	return time.Unix(0, newTimeOfNextPermissionIssue)
+	return time.Unix(0, newTimeOfNextPermissionIssue), time.Duration(newTimeOfNextPermissionIssue - now)
 }

limiter_mutexbased.go

@@ -51,6 +51,19 @@ func newMutexBased(rate int, opts ...Option) *mutexLimiter {
 // Take blocks to ensure that the time spent between multiple
 // Take calls is on average per/rate.
 func (t *mutexLimiter) Take() time.Time {
+	last, interval := t.take()
+	t.clock.Sleep(interval)
+
+	return last
+}
+
+func (t *mutexLimiter) TakeNoWait() (time.Time, time.Duration) {
+	return t.take()
+}
+
+func (t *mutexLimiter) take() (time.Time, time.Duration) {
+	interval := time.Duration(0)
+
 	t.Lock()
 	defer t.Unlock()
 
@@ -59,7 +72,7 @@ func (t *mutexLimiter) Take() time.Time {
 	// If this is our first request, then we allow it.
 	if t.last.IsZero() {
 		t.last = now
-		return t.last
+		return t.last, interval
 	}
 
 	// sleepFor calculates how much time we should sleep based on
@@ -77,12 +90,12 @@ func (t *mutexLimiter) Take() time.Time {
 
 	// If sleepFor is positive, then we should sleep now.
 	if t.sleepFor > 0 {
-		t.clock.Sleep(t.sleepFor)
+		interval = t.sleepFor
 		t.last = now.Add(t.sleepFor)
 		t.sleepFor = 0
 	} else {
 		t.last = now
 	}
 
-	return t.last
+	return t.last, interval
 }

ratelimit.go

@@ -35,6 +35,7 @@ import (
 type Limiter interface {
 	// Take should block to make sure that the RPS is met.
 	Take() time.Time
+	TakeNoWait() (time.Time, time.Duration)
 }
 
 // Clock is the minimum necessary interface to instantiate a rate limiter with
@@ -133,3 +134,7 @@ func NewUnlimited() Limiter {
 func (unlimited) Take() time.Time {
 	return time.Now()
 }
+
+func (unlimited) TakeNoWait() (time.Time, time.Duration) {
+	return time.Now(), time.Duration(0)
+}