All old expressed in terms of new

Author: nelz9999

ratelimit/token_bucket.go

@@ -23,16 +23,11 @@ func NewTokenBucketLimiter(tb *ratelimit.Bucket) endpoint.Middleware {
 
 // NewTokenBucketThrottler returns an endpoint.Middleware that acts as a
 // request throttler based on a token-bucket algorithm. Requests that would
-// exceed the maximum request rate are delayed via the parameterized sleep
-// function. By default you may pass time.Sleep.
-func NewTokenBucketThrottler(tb *ratelimit.Bucket, sleep func(time.Duration)) endpoint.Middleware {
-	// return NewDelayingLimiter(NewWaiter(tb))
-	return func(next endpoint.Endpoint) endpoint.Endpoint {
-		return func(ctx context.Context, request interface{}) (interface{}, error) {
-			sleep(tb.Take(1))
-			return next(ctx, request)
-		}
-	}
+// exceed the maximum request rate are delayed.
+// The parameterized function "_" is kept for backwards-compatiblity of
+// the API, but it is no longer used for anything. You may pass it nil.
+func NewTokenBucketThrottler(tb *ratelimit.Bucket, _ func(time.Duration)) endpoint.Middleware {
+	return NewDelayingLimiter(NewWaiter(tb))
 }
 
 // Allower dictates whether or not a request is acceptable to run.

ratelimit/token_bucket_test.go

@@ -2,82 +2,62 @@ package ratelimit_test
 
 import (
 	"context"
-	"math"
 	"strings"
 	"testing"
 	"time"
 
 	jujuratelimit "github.com/juju/ratelimit"
+	"golang.org/x/time/rate"
 
 	"github.com/go-kit/kit/endpoint"
 	"github.com/go-kit/kit/ratelimit"
-	"golang.org/x/time/rate"
 )
 
+var nopEndpoint = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
+
 func TestTokenBucketLimiter(t *testing.T) {
-	e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
-	for _, n := range []int{1, 2, 100} {
-		tb := jujuratelimit.NewBucketWithRate(float64(n), int64(n))
-		testLimiter(t, ratelimit.NewTokenBucketLimiter(tb)(e), n)
-	}
+	tb := jujuratelimit.NewBucket(time.Minute, 1)
+	testSuccessThenFailure(
+		t,
+		ratelimit.NewTokenBucketLimiter(tb)(nopEndpoint),
+		ratelimit.ErrLimited.Error())
 }
 
 func TestTokenBucketThrottler(t *testing.T) {
-	d := time.Duration(0)
-	s := func(d0 time.Duration) { d = d0 }
-
-	e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
-	e = ratelimit.NewTokenBucketThrottler(jujuratelimit.NewBucketWithRate(1, 1), s)(e)
-
-	// First request should go through with no delay.
-	e(context.Background(), struct{}{})
-	if want, have := time.Duration(0), d; want != have {
-		t.Errorf("want %s, have %s", want, have)
-	}
-
-	// Next request should request a ~1s sleep.
-	e(context.Background(), struct{}{})
-	if want, have, tol := time.Second, d, time.Millisecond; math.Abs(float64(want-have)) > float64(tol) {
-		t.Errorf("want %s, have %s", want, have)
-	}
-}
-
-func testLimiter(t *testing.T, e endpoint.Endpoint, rate int) {
-	// First <rate> requests should succeed.
-	for i := 0; i < rate; i++ {
-		if _, err := e(context.Background(), struct{}{}); err != nil {
-			t.Fatalf("rate=%d: request %d/%d failed: %v", rate, i+1, rate, err)
-		}
-	}
-
-	// Next request should fail.
-	if _, err := e(context.Background(), struct{}{}); err != ratelimit.ErrLimited {
-		t.Errorf("rate=%d: want %v, have %v", rate, ratelimit.ErrLimited, err)
-	}
+	tb := jujuratelimit.NewBucket(time.Minute, 1)
+	testSuccessThenFailure(
+		t,
+		ratelimit.NewTokenBucketThrottler(tb, nil)(nopEndpoint),
+		"context deadline exceeded")
 }
 
 func TestXRateErroring(t *testing.T) {
 	limit := rate.NewLimiter(rate.Every(time.Minute), 1)
-	e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
-	testLimiter(t, ratelimit.NewErroringLimiter(limit)(e), 1)
+	testSuccessThenFailure(
+		t,
+		ratelimit.NewErroringLimiter(limit)(nopEndpoint),
+		ratelimit.ErrLimited.Error())
 }
 
 func TestXRateDelaying(t *testing.T) {
 	limit := rate.NewLimiter(rate.Every(time.Minute), 1)
-	e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
-	e = ratelimit.NewDelayingLimiter(limit)(e)
+	testSuccessThenFailure(
+		t,
+		ratelimit.NewDelayingLimiter(limit)(nopEndpoint),
+		"exceed context deadline")
+}
 
-	_, err := e(context.Background(), struct{}{})
-	if err != nil {
+func testSuccessThenFailure(t *testing.T, e endpoint.Endpoint, failContains string) {
+	ctx, cxl := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cxl()
+
+	// First request should succeed.
+	if _, err := e(ctx, struct{}{}); err != nil {
 		t.Errorf("unexpected: %v\n", err)
 	}
 
-	dur := 500 * time.Millisecond
-	ctx, cxl := context.WithTimeout(context.Background(), dur)
-	defer cxl()
-
-	_, err = e(ctx, struct{}{})
-	if !strings.Contains(err.Error(), "exceed context deadline") {
-		t.Errorf("expected timeout: %v\n", err)
+	// Next request should fail.
+	if _, err := e(ctx, struct{}{}); !strings.Contains(err.Error(), failContains) {
+		t.Errorf("expected `%s`: %v\n", failContains, err)
 	}
 }