diff --git a/ratelimit/token_bucket.go b/ratelimit/token_bucket.go index b71e50b..e8a6de6 100644 --- a/ratelimit/token_bucket.go +++ b/ratelimit/token_bucket.go @@ -3,9 +3,6 @@ package ratelimit import ( "context" "errors" - "time" - - "github.com/juju/ratelimit" "github.com/go-kit/kit/endpoint" ) @@ -14,22 +11,6 @@ import ( // triggered and the request is rejected. var ErrLimited = errors.New("rate limit exceeded") -// NewTokenBucketLimiter returns an endpoint.Middleware that acts as a rate -// limiter based on a token-bucket algorithm. Requests that would exceed the -// maximum request rate are simply rejected with an error. -func NewTokenBucketLimiter(tb *ratelimit.Bucket) endpoint.Middleware { - return NewErroringLimiter(NewAllower(tb)) -} - -// 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. -// 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. // The Limiter from "golang.org/x/time/rate" already implements this interface, // one is able to use that in NewErroringLimiter without any modifications. @@ -81,13 +62,6 @@ func (f AllowerFunc) Allow() bool { return f() } -// NewAllower turns an existing ratelimit.Bucket into an API-compatible form -func NewAllower(tb *ratelimit.Bucket) Allower { - return AllowerFunc(func() bool { - return (tb.TakeAvailable(1) != 0) - }) -} - // WaiterFunc is an adapter that lets a function operate as if // it implements Waiter type WaiterFunc func(ctx context.Context) error @@ -96,17 +70,3 @@ type WaiterFunc func(ctx context.Context) error func (f WaiterFunc) Wait(ctx context.Context) error { return f(ctx) } - -// NewWaiter turns an existing ratelimit.Bucket into an API-compatible form -func NewWaiter(tb *ratelimit.Bucket) Waiter { - return WaiterFunc(func(ctx context.Context) error { - dur := tb.Take(1) - select { - case <-ctx.Done(): - return ctx.Err() - case <-time.After(dur): - // happy path - } - return nil - }) -} diff --git a/ratelimit/token_bucket_test.go b/ratelimit/token_bucket_test.go index d444fe9..3845c9e 100644 --- a/ratelimit/token_bucket_test.go +++ b/ratelimit/token_bucket_test.go @@ -6,7 +6,6 @@ import ( "testing" "time" - jujuratelimit "github.com/juju/ratelimit" "golang.org/x/time/rate" "github.com/go-kit/kit/endpoint" @@ -15,22 +14,6 @@ import ( var nopEndpoint = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil } -func TestTokenBucketLimiter(t *testing.T) { - tb := jujuratelimit.NewBucket(time.Minute, 1) - testSuccessThenFailure( - t, - ratelimit.NewTokenBucketLimiter(tb)(nopEndpoint), - ratelimit.ErrLimited.Error()) -} - -func TestTokenBucketThrottler(t *testing.T) { - 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) testSuccessThenFailure(