Ratelimit interceptor

ratelimit/doc.go

@@ -0,0 +1,14 @@
+// See LICENSE for licensing terms.
+
+/*
+`ratelimit` a generic server-side ratelimit middleware for gRPC.
+
+Server Side Ratelimit Middleware
+
+It allows to do grpc rate limit by your own rate limiter (e.g. token bucket, leaky bucket, etc.)
+
+`ratelimit/tokenbucket`provides an implementation based on token bucket `github.com/juju/ratelimit`.
+
+Please see examples for simple examples of use.
+*/
+package ratelimit

ratelimit/examples_test.go

@@ -0,0 +1,33 @@
+// Copyright 2018 Zheng Dayu. All Rights Reserved.
+// See LICENSE for licensing terms.
+
+package ratelimit_test
+
+import (
+ "time"
+
+ "github.com/grpc-ecosystem/go-grpc-middleware"
+ "github.com/grpc-ecosystem/go-grpc-middleware/ratelimit"
+ "github.com/grpc-ecosystem/go-grpc-middleware/ratelimit/tokenbucket"
+ "google.golang.org/grpc"
+)
+
+// Simple example of server initialization code.
+func Example() {
+ // Create unary/stream rateLimiters, based on token bucket here.
+ // You can implement your own ratelimiter for the interface.
+ unaryRateLimiter := tokenbucket.NewTokenBucketRateLimiter(1*time.Second, 10, 10, 10*time.Second)
+ streamRateLimiter := tokenbucket.NewTokenBucketRateLimiter(1*time.Second, 5, 5, 5*time.Second)
+ _ = grpc.NewServer(
+ grpc_middleware.WithUnaryServerChain(
+ ratelimit.UnaryServerInterceptor(
+ ratelimit.WithRateLimiter(unaryRateLimiter),
+ ),
+ ),
+ grpc_middleware.WithStreamServerChain(
+ ratelimit.StreamServerInterceptor(
+ ratelimit.WithRateLimiter(streamRateLimiter),
+ ),
+ ),
+ )
+}

ratelimit/options.go

@@ -0,0 +1,16 @@
+// See LICENSE for licensing terms.
+
+package ratelimit
+
+import "time"
+
+const infinityDuration time.Duration = 0x7fffffffffffffff
+
+type Option func(*rateLimiter)
+
+// WithRateLimiter customizes your limiter in the middleware
+func WithRateLimiter(l Limiter) Option {
+ return func(r *rateLimiter) {
+ r.limiter = l
+ }
+}

ratelimit/ratelimit.go

@@ -0,0 +1,63 @@
+// See LICENSE for licensing terms.
+
+package ratelimit
+
+import (
+ "fmt"
+
+ "golang.org/x/net/context"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/codes"
+ "google.golang.org/grpc/status"
+)
+
+type Limiter interface {
+ Limit() bool
+}
+
+type rateLimiter struct {
+ limiter Limiter
+}
+
+type emptyLimiter struct{}
+
+func (e *emptyLimiter) Limit() bool {
+ return false
+}
+
+func emptyRatelimiter() *rateLimiter {
+ return &rateLimiter{
+ limiter: &emptyLimiter{},
+ }
+}
+
+// UnaryServerInterceptor returns a new unary server interceptors that performs request rate limit.
+func UnaryServerInterceptor(opts ...Option) grpc.UnaryServerInterceptor {
+ ratelimiter := emptyRatelimiter()
+ for _, opt := range opts {
+ opt(ratelimiter)
+ }
+ return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+ if ratelimiter.limiter.Limit() {
+ return nil, status.Errorf(codes.ResourceExhausted, "%s is rejected by grpc_ratelimit middleare, please retry later.", info.FullMethod)
+ }
+ return handler(ctx, req)
+ }
+}
+
+// StreamServerInterceptor returns a new stream server interceptor that performs rate limiting on the request.
+func StreamServerInterceptor(opts ...Option) grpc.StreamServerInterceptor {
+ ratelimiter := emptyRatelimiter()
+ fmt.Println(ratelimiter.limiter.Limit())
+ for _, opt := range opts {
+ opt(ratelimiter)
+ }
+ fmt.Println(ratelimiter.limiter.Limit())
+ return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+ fmt.Println(ratelimiter.limiter.Limit())
+ if ratelimiter.limiter.Limit() {
+ return status.Errorf(codes.ResourceExhausted, "%s is rejected by grpc_ratelimit middleare, please retry later.", info.FullMethod)
+ }
+ return handler(srv, stream)
+ }
+}

ratelimit/ratelimit_test.go

@@ -0,0 +1,106 @@
+package ratelimit
+
+import (
+ "context"
+ "errors"
+ "testing"
+
+ "google.golang.org/grpc"
+
+ "github.com/stretchr/testify/assert"
+)
+
+const errMsgFake = "fake error"
+
+func TestUnaryServerInterceptor_NoLimit(t *testing.T) {
+ interceptor := UnaryServerInterceptor()
+ handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+ return nil, errors.New(errMsgFake)
+ }
+ req, err := interceptor(nil, nil, nil, handler)
+ assert.Nil(t, req)
+ assert.EqualError(t, err, errMsgFake)
+}
+
+type mockPassLimiter struct{}
+
+func (*mockPassLimiter) Limit() bool {
+ return false
+}
+
+func TestUnaryServerInterceptor_RateLimitPass(t *testing.T) {
+ unaryRateLimiter := &mockPassLimiter{}
+ interceptor := UnaryServerInterceptor(
+ WithRateLimiter(unaryRateLimiter),
+ )
+ handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+ return nil, errors.New(errMsgFake)
+ }
+ info := &grpc.UnaryServerInfo{
+ FullMethod: "FakeMethod",
+ }
+ req, err := interceptor(nil, nil, info, handler)
+ assert.Nil(t, req)
+ assert.EqualError(t, err, errMsgFake)
+}
+
+type mockFailLimiter struct{}
+
+func (*mockFailLimiter) Limit() bool {
+ return true
+}
+
+func TestUnaryServerInterceptor_RateLimitFail(t *testing.T) {
+ unaryRateLimiter := &mockFailLimiter{}
+ interceptor := UnaryServerInterceptor(
+ WithRateLimiter(unaryRateLimiter),
+ )
+ handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+ return nil, errors.New(errMsgFake)
+ }
+ info := &grpc.UnaryServerInfo{
+ FullMethod: "FakeMethod",
+ }
+ req, err := interceptor(nil, nil, info, handler)
+ assert.Nil(t, req)
+ assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = FakeMethod is rejected by grpc_ratelimit middleare, please retry later.")
+}
+
+func TestStreamServerInterceptor_NoLimit(t *testing.T) {
+ interceptor := StreamServerInterceptor()
+ handler := func(srv interface{}, stream grpc.ServerStream) error {
+ return errors.New(errMsgFake)
+ }
+ err := interceptor(nil, nil, nil, handler)
+ assert.EqualError(t, err, errMsgFake)
+}
+
+func TestStreamServerInterceptor_RateLimitPass(t *testing.T) {
+ streamRateLimiter := &mockPassLimiter{}
+ interceptor := StreamServerInterceptor(
+ WithRateLimiter(streamRateLimiter),
+ )
+ handler := func(srv interface{}, stream grpc.ServerStream) error {
+ return errors.New(errMsgFake)
+ }
+ info := &grpc.StreamServerInfo{
+ FullMethod: "FakeMethod",
+ }
+ err := interceptor(nil, nil, info, handler)
+ assert.EqualError(t, err, errMsgFake)
+}
+
+func TestStreamServerInterceptor_RateLimitFail(t *testing.T) {
+ streamRateLimiter := &mockFailLimiter{}
+ interceptor := StreamServerInterceptor(
+ WithRateLimiter(streamRateLimiter),
+ )
+ handler := func(srv interface{}, stream grpc.ServerStream) error {
+ return errors.New(errMsgFake)
+ }
+ info := &grpc.StreamServerInfo{
+ FullMethod: "FakeMethod",
+ }
+ err := interceptor(nil, nil, info, handler)
+ assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = FakeMethod is rejected by grpc_ratelimit middleare, please retry later.")
+}

ratelimit/tokenbucket/tokenbucket.go

@@ -0,0 +1,28 @@
+// Copyright 2018 Zheng Dayu. All Rights Reserved.
+// See LICENSE for licensing terms.
+
+package tokenbucket
+
+import (
+ "time"
+
+ "github.com/juju/ratelimit"
+)
+
+type tokenBucketLimiter struct {
+ limiter *ratelimit.Bucket
+ maxWaitDuration time.Duration
+}
+
+// NewTokenBucketRateLimiter creates a tokenBucketLimiter.
+func NewTokenBucketRateLimiter(fillInterval time.Duration, capacity, quantum int64, maxWaitDuration time.Duration) *tokenBucketLimiter {
+ return &tokenBucketLimiter{
+ limiter: ratelimit.NewBucketWithQuantum(fillInterval, capacity, quantum),
+ maxWaitDuration: maxWaitDuration,
+ }
+}
+
+// Limit takes 1 token from the bucket only if it needs to wait for no greater than maxWaitDuration
+func (b *tokenBucketLimiter) Limit() bool {
+ return !b.limiter.WaitMaxDuration(1, b.maxWaitDuration)
+}

ratelimit/tokenbucket/tokenbucket_test.go

@@ -0,0 +1,22 @@
+package tokenbucket
+
+import (
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+)
+
+func TestTokenBucketRateLimiter_LimitPass(t *testing.T) {
+ l := NewTokenBucketRateLimiter(1*time.Millisecond, 1, 1, 1*time.Millisecond)
+ ok := l.Limit()
+ assert.False(t, ok)
+}
+
+func TestTokenBucketRateLimiter_LimitFail(t *testing.T) {
+ l := NewTokenBucketRateLimiter(10*time.Second, 1, 1, 1*time.Millisecond)
+ ok := l.Limit()
+ assert.False(t, ok)
+ ok2 := l.Limit()
+ assert.True(t, ok2)
+}