Add ratelimit

endpoint/ratelimit/token_bucket.go

@@ -0,0 +1,72 @@
+package ratelimit
+
+import (
+	"context"
+	"errors"
+
+	"github.com/dreamsxin/go-kit/endpoint"
+)
+
+// ErrLimited is returned in the request path when the rate limiter is
+// triggered and the request is rejected.
+var ErrLimited = errors.New("rate limit exceeded")
+
+// 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.
+type Allower interface {
+	Allow() bool
+}
+
+// NewErroringLimiter returns an endpoint.Middleware that acts as a rate
+// limiter. Requests that would exceed the
+// maximum request rate are simply rejected with an error.
+func NewErroringLimiter(limit Allower) endpoint.Middleware {
+	return func(next endpoint.Endpoint) endpoint.Endpoint {
+		return func(ctx context.Context, request interface{}) (interface{}, error) {
+			if !limit.Allow() {
+				return nil, ErrLimited
+			}
+			return next(ctx, request)
+		}
+	}
+}
+
+// Waiter dictates how long a request must be delayed.
+// The Limiter from "golang.org/x/time/rate" already implements this interface,
+// one is able to use that in NewDelayingLimiter without any modifications.
+type Waiter interface {
+	Wait(ctx context.Context) error
+}
+
+// NewDelayingLimiter returns an endpoint.Middleware that acts as a
+// request throttler. Requests that would
+// exceed the maximum request rate are delayed via the Waiter function
+func NewDelayingLimiter(limit Waiter) endpoint.Middleware {
+	return func(next endpoint.Endpoint) endpoint.Endpoint {
+		return func(ctx context.Context, request interface{}) (interface{}, error) {
+			if err := limit.Wait(ctx); err != nil {
+				return nil, err
+			}
+			return next(ctx, request)
+		}
+	}
+}
+
+// AllowerFunc is an adapter that lets a function operate as if
+// it implements Allower
+type AllowerFunc func() bool
+
+// Allow makes the adapter implement Allower
+func (f AllowerFunc) Allow() bool {
+	return f()
+}
+
+// WaiterFunc is an adapter that lets a function operate as if
+// it implements Waiter
+type WaiterFunc func(ctx context.Context) error
+
+// Wait makes the adapter implement Waiter
+func (f WaiterFunc) Wait(ctx context.Context) error {
+	return f(ctx)
+}

endpoint/ratelimit/token_bucket_test.go

@@ -0,0 +1,46 @@
+package ratelimit_test
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"golang.org/x/time/rate"
+
+	"github.com/dreamsxin/go-kit/endpoint"
+	"github.com/dreamsxin/go-kit/endpoint/ratelimit"
+)
+
+var nopEndpoint = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
+
+func TestXRateErroring(t *testing.T) {
+	limit := rate.NewLimiter(rate.Every(time.Minute), 1)
+	testSuccessThenFailure(
+		t,
+		ratelimit.NewErroringLimiter(limit)(nopEndpoint),
+		ratelimit.ErrLimited.Error())
+}
+
+func TestXRateDelaying(t *testing.T) {
+	limit := rate.NewLimiter(rate.Every(time.Minute), 1)
+	testSuccessThenFailure(
+		t,
+		ratelimit.NewDelayingLimiter(limit)(nopEndpoint),
+		"exceed context deadline")
+}
+
+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)
+	}
+
+	// Next request should fail.
+	if _, err := e(ctx, struct{}{}); !strings.Contains(err.Error(), failContains) {
+		t.Errorf("expected `%s`: %v\n", failContains, err)
+	}
+}

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/sony/gobreaker v1.0.0
 	github.com/streadway/handy v0.0.0-20200128134331-0f66f006fb2e
 	go.uber.org/zap v1.26.0
+	golang.org/x/time v0.5.0
 	google.golang.org/grpc v1.61.0
 	google.golang.org/protobuf v1.31.0
 )

go.sum

@@ -235,6 +235,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
+golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
+golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190907020128-2ca718005c18/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=