Update the Limiter interface to accept a context.Context object

ratelimit.go

@@ -24,7 +24,8 @@ import (
 	"sync"
 	"time"
 
-	"go.uber.org/ratelimit/internal/clock"
+	"context"
+	"math"
 )
 
 // Note: This file is inspired by:
@@ -35,24 +36,16 @@ import (
 // may block to throttle the goroutine.
 type Limiter interface {
 	// Take should block to make sure that the RPS is met.
-	Take() time.Time
-}
-
-// Clock is the minimum necessary interface to instantiate a rate limiter with
-// a clock or mock clock, compatible with clocks created using
-// github.com/andres-erbsen/clock.
-type Clock interface {
-	Now() time.Time
-	Sleep(time.Duration)
+	Take(context.Context) bool
 }
 
 type limiter struct {
 	sync.Mutex
 	last       time.Time
+	timer      *time.Timer
 	sleepFor   time.Duration
 	perRequest time.Duration
 	maxSlack   time.Duration
-	clock      Clock
 }
 
 // Option configures a Limiter.
@@ -63,24 +56,14 @@ func New(rate int, opts ...Option) Limiter {
 	l := &limiter{
 		perRequest: time.Second / time.Duration(rate),
 		maxSlack:   -10 * time.Second / time.Duration(rate),
+		timer:      time.NewTimer(time.Duration(math.MaxInt64)),
 	}
 	for _, opt := range opts {
 		opt(l)
 	}
-	if l.clock == nil {
-		l.clock = clock.New()
-	}
 	return l
 }
 
-// WithClock returns an option for ratelimit.New that provides an alternate
-// Clock implementation, typically a mock Clock for testing.
-func WithClock(clock Clock) Option {
-	return func(l *limiter) {
-		l.clock = clock
-	}
-}
-
 // WithoutSlack is an option for ratelimit.New that initializes the limiter
 // without any initial tolerance for bursts of traffic.
 var WithoutSlack Option = withoutSlackOption
@@ -91,23 +74,24 @@ func withoutSlackOption(l *limiter) {
 
 // Take blocks to ensure that the time spent between multiple
 // Take calls is on average time.Second/rate.
-func (t *limiter) Take() time.Time {
+func (t *limiter) Take(ctx context.Context) bool {
 	t.Lock()
 	defer t.Unlock()
 
-	now := t.clock.Now()
+	now := time.Now()
 
 	// If this is our first request, then we allow it.
 	if t.last.IsZero() {
 		t.last = now
-		return t.last
+		return true
 	}
 
 	// sleepFor calculates how much time we should sleep based on
 	// the perRequest budget and how long the last request took.
 	// Since the request may take longer than the budget, this number
 	// can get negative, and is summed across requests.
 	t.sleepFor += t.perRequest - now.Sub(t.last)
+	t.last = now
 
 	// We shouldn't allow sleepFor to get too negative, since it would mean that
 	// a service that slowed down a lot for a short period of time would get
@@ -118,14 +102,21 @@ func (t *limiter) Take() time.Time {
 
 	// If sleepFor is positive, then we should sleep now.
 	if t.sleepFor > 0 {
-		t.clock.Sleep(t.sleepFor)
+		if !t.timer.Stop() {
+			<-t.timer.C
+		}
+		t.timer.Reset(t.sleepFor)
+		select {
+		case <-t.timer.C:
+		case <-ctx.Done():
+			return false
+		}
+
 		t.last = now.Add(t.sleepFor)
 		t.sleepFor = 0
-	} else {
-		t.last = now
 	}
 
-	return t.last
+	return true
 }
 
 type unlimited struct{}
@@ -135,6 +126,6 @@ func NewUnlimited() Limiter {
 	return unlimited{}
 }
 
-func (unlimited) Take() time.Time {
-	return time.Now()
+func (unlimited) Take(_ context.Context) bool {
+	return true
 }

ratelimit_test.go

@@ -5,9 +5,9 @@ import (
 	"sync"
 	"testing"
 	"time"
+	"context"
 
 	"go.uber.org/ratelimit"
-	"go.uber.org/ratelimit/internal/clock"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/uber-go/atomic"
@@ -17,8 +17,11 @@ func ExampleRatelimit() {
 	rl := ratelimit.New(100) // per second
 
 	prev := time.Now()
+	ctx := context.Background()
 	for i := 0; i < 10; i++ {
-		now := rl.Take()
+		rl.Take(ctx)
+		now := time.Now()
+
 		if i > 0 {
 			fmt.Println(i, now.Sub(prev))
 		}
@@ -40,8 +43,9 @@ func ExampleRatelimit() {
 func TestUnlimited(t *testing.T) {
 	now := time.Now()
 	rl := ratelimit.NewUnlimited()
+	ctx := context.Background()
 	for i := 0; i < 1000; i++ {
-		rl.Take()
+		rl.Take(ctx)
 	}
 	assert.Condition(t, func() bool { return time.Now().Sub(now) < 1*time.Millisecond }, "no artificial delay")
 }
@@ -51,91 +55,81 @@ func TestRateLimiter(t *testing.T) {
 	wg.Add(1)
 	defer wg.Wait()
 
-	clock := clock.NewMock()
-	rl := ratelimit.New(100, ratelimit.WithClock(clock), ratelimit.WithoutSlack)
+	rl := ratelimit.New(100, ratelimit.WithoutSlack)
 
 	count := atomic.NewInt32(0)
 
 	// Until we're done...
-	done := make(chan struct{})
-	defer close(done)
+	ctx := context.Background()
 
 	// Create copious counts concurrently.
-	go job(rl, count, done)
-	go job(rl, count, done)
-	go job(rl, count, done)
-	go job(rl, count, done)
+	go job(rl, count, ctx)
+	go job(rl, count, ctx)
+	go job(rl, count, ctx)
+	go job(rl, count, ctx)
 
-	clock.AfterFunc(1*time.Second, func() {
+	time.AfterFunc(1*time.Second, func() {
 		assert.InDelta(t, 100, count.Load(), 10, "count within rate limit")
 	})
 
-	clock.AfterFunc(2*time.Second, func() {
+	time.AfterFunc(2*time.Second, func() {
 		assert.InDelta(t, 200, count.Load(), 10, "count within rate limit")
 	})
 
-	clock.AfterFunc(3*time.Second, func() {
+	time.AfterFunc(3*time.Second, func() {
 		assert.InDelta(t, 300, count.Load(), 10, "count within rate limit")
 		wg.Done()
 	})
-
-	clock.Add(4 * time.Second)
-
-	clock.Add(5 * time.Second)
 }
 
 func TestDelayedRateLimiter(t *testing.T) {
 	var wg sync.WaitGroup
 	wg.Add(1)
 	defer wg.Wait()
 
-	clock := clock.NewMock()
-	slow := ratelimit.New(10, ratelimit.WithClock(clock))
-	fast := ratelimit.New(100, ratelimit.WithClock(clock))
+	slow := ratelimit.New(10)
+	fast := ratelimit.New(100)
 
 	count := atomic.NewInt32(0)
 
 	// Until we're done...
-	done := make(chan struct{})
-	defer close(done)
+	ctx := context.Background()
 
 	// Run a slow job
 	go func() {
 		for {
-			slow.Take()
-			fast.Take()
+			slow.Take(ctx)
+			fast.Take(ctx)
 			count.Inc()
 			select {
-			case <-done:
+			case <-ctx.Done():
 				return
 			default:
 			}
 		}
 	}()
 
 	// Accumulate slack for 10 seconds,
-	clock.AfterFunc(20*time.Second, func() {
+	time.AfterFunc(20*time.Second, func() {
 		// Then start working.
-		go job(fast, count, done)
-		go job(fast, count, done)
-		go job(fast, count, done)
-		go job(fast, count, done)
+		go job(fast, count, ctx)
+		go job(fast, count, ctx)
+		go job(fast, count, ctx)
+		go job(fast, count, ctx)
 	})
 
-	clock.AfterFunc(30*time.Second, func() {
+	time.AfterFunc(30*time.Second, func() {
 		assert.InDelta(t, 1200, count.Load(), 10, "count within rate limit")
 		wg.Done()
 	})
-
-	clock.Add(40 * time.Second)
 }
 
-func job(rl ratelimit.Limiter, count *atomic.Int32, done <-chan struct{}) {
+func job(rl ratelimit.Limiter, count *atomic.Int32, ctx context.Context) {
 	for {
-		rl.Take()
+		rl.Take(ctx)
 		count.Inc()
 		select {
-		case <-done:
+		case <-ctx.Done():
 			return
 		default:
 		}