mirror of
https://github.com/go-kit/kit.git
synced 2025-07-15 01:04:44 +02:00
Remove dependency on juju
This commit is contained in:
@ -3,9 +3,6 @@ package ratelimit
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
"time"
|
|
||||||
|
|
||||||
"github.com/juju/ratelimit"
|
|
||||||
|
|
||||||
"github.com/go-kit/kit/endpoint"
|
"github.com/go-kit/kit/endpoint"
|
||||||
)
|
)
|
||||||
@ -14,22 +11,6 @@ import (
|
|||||||
// triggered and the request is rejected.
|
// triggered and the request is rejected.
|
||||||
var ErrLimited = errors.New("rate limit exceeded")
|
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.
|
// Allower dictates whether or not a request is acceptable to run.
|
||||||
// The Limiter from "golang.org/x/time/rate" already implements this interface,
|
// The Limiter from "golang.org/x/time/rate" already implements this interface,
|
||||||
// one is able to use that in NewErroringLimiter without any modifications.
|
// one is able to use that in NewErroringLimiter without any modifications.
|
||||||
@ -81,13 +62,6 @@ func (f AllowerFunc) Allow() bool {
|
|||||||
return f()
|
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
|
// WaiterFunc is an adapter that lets a function operate as if
|
||||||
// it implements Waiter
|
// it implements Waiter
|
||||||
type WaiterFunc func(ctx context.Context) error
|
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 {
|
func (f WaiterFunc) Wait(ctx context.Context) error {
|
||||||
return f(ctx)
|
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
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
@ -6,7 +6,6 @@ import (
|
|||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
jujuratelimit "github.com/juju/ratelimit"
|
|
||||||
"golang.org/x/time/rate"
|
"golang.org/x/time/rate"
|
||||||
|
|
||||||
"github.com/go-kit/kit/endpoint"
|
"github.com/go-kit/kit/endpoint"
|
||||||
@ -15,22 +14,6 @@ import (
|
|||||||
|
|
||||||
var nopEndpoint = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
|
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) {
|
func TestXRateErroring(t *testing.T) {
|
||||||
limit := rate.NewLimiter(rate.Every(time.Minute), 1)
|
limit := rate.NewLimiter(rate.Every(time.Minute), 1)
|
||||||
testSuccessThenFailure(
|
testSuccessThenFailure(
|
||||||
|
Reference in New Issue
Block a user