Ratelimit interceptor

ratelimit/DOC.md

@@ -0,0 +1,105 @@
+# grpc_ratelimit
+`import "github.com/grpc-ecosystem/go-grpc-middleware/ratelimit"`
+
+* [Overview](#pkg-overview)
+* [Imported Packages](#pkg-imports)
+* [Index](#pkg-index)
+* [Examples](#pkg-examples)
+
+## <a name="pkg-overview">Overview</a>
+`grpc_ratelimit` a generic server-side ratelimit middleware for gRPC.
+
+### Server Side Ratelimit Middleware
+It allows to use your own rate limiter (e.g. token bucket, leaky bucket, etc.) to do grpc rate limit.
+
+`ratelimit/tokenbucket`provides an implementation based on token bucket `github.com/juju/ratelimit`.
+
+Please see examples for simple examples of use.
+
+#### Example:
+
+<details>
+<summary>Click to expand code.</summary>
+
+```go
+// 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)
+streamRateLimiter := tokenbucket.NewTokenBucketRateLimiter(1*time.Second, 5, 5)
+_ = grpc.NewServer(
+ grpc_middleware.WithUnaryServerChain(
+ grpc_ratelimit.UnaryServerInterceptor(
+ grpc_ratelimit.WithLimiter(unaryRateLimiter),
+ grpc_ratelimit.WithMaxWaitDuration(10*time.Second),
+ ),
+ ),
+ grpc_middleware.WithStreamServerChain(
+ grpc_ratelimit.StreamServerInterceptor(
+ grpc_ratelimit.WithLimiter(streamRateLimiter),
+ grpc_ratelimit.WithMaxWaitDuration(5*time.Second),
+ ),
+ ),
+)
+```
+
+</details>
+
+## <a name="pkg-imports">Imported Packages</a>
+
+- [golang.org/x/net/context](https://godoc.org/golang.org/x/net/context)
+- [google.golang.org/grpc](https://godoc.org/google.golang.org/grpc)
+- [google.golang.org/grpc/codes](https://godoc.org/google.golang.org/grpc/codes)
+- [google.golang.org/grpc/status](https://godoc.org/google.golang.org/grpc/status)
+
+## <a name="pkg-index">Index</a>
+* [func StreamServerInterceptor(opts ...Option) grpc.StreamServerInterceptor](#StreamServerInterceptor)
+* [func UnaryServerInterceptor(opts ...Option) grpc.UnaryServerInterceptor](#UnaryServerInterceptor)
+* [type Limiter](#Limiter)
+* [type Option](#Option)
+ * [func WithLimiter(l Limiter) Option](#WithLimiter)
+ * [func WithMaxWaitDuration(maxWaitDuration time.Duration) Option](#WithMaxWaitDuration)
+
+#### <a name="pkg-examples">Examples</a>
+* [Package (Initialization)](#example__initialization)
+
+#### <a name="pkg-files">Package files</a>
+[doc.go](./doc.go) [options.go](./options.go) [ratelimit.go](./ratelimit.go) 
+
+## <a name="StreamServerInterceptor">func</a> [StreamServerInterceptor](./ratelimit.go#L55)
+``` go
+func StreamServerInterceptor(opts ...Option) grpc.StreamServerInterceptor
+```
+StreamServerInterceptor returns a new stream server interceptors that performs request rate limit.
+
+## <a name="UnaryServerInterceptor">func</a> [UnaryServerInterceptor](./ratelimit.go#L41)
+``` go
+func UnaryServerInterceptor(opts ...Option) grpc.UnaryServerInterceptor
+```
+UnaryServerInterceptor returns a new unary server interceptors that performs request rate limit.
+
+## <a name="Limiter">type</a> [Limiter](./ratelimit.go#L15-L17)
+``` go
+type Limiter interface {
+ WaitMaxDuration(time.Duration) bool
+}
+```
+
+## <a name="Option">type</a> [Option](./options.go#L10)
+``` go
+type Option func(*rateLimiter)
+```
+
+### <a name="WithLimiter">func</a> [WithLimiter](./options.go#L13)
+``` go
+func WithLimiter(l Limiter) Option
+```
+WithLimiter customizes your limiter in the middleware
+
+### <a name="WithMaxWaitDuration">func</a> [WithMaxWaitDuration](./options.go#L23)
+``` go
+func WithMaxWaitDuration(maxWaitDuration time.Duration) Option
+```
+WithMaxWaitDuration customizes maxWaitDuration in limiter's WaitMaxDuration action.
+
+- - -
+Generated by [godoc2ghmd](https://github.com/GandalfUK/godoc2ghmd)

ratelimit/README.md

@@ -0,0 +1 @@
+DOC.md

ratelimit/doc.go

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

ratelimit/examples_test.go

@@ -0,0 +1,35 @@
+// Copyright 2018 Zheng Dayu. All Rights Reserved.
+// See LICENSE for licensing terms.
+
+package grpc_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_initialization() {
+ // 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)
+ streamRateLimiter := tokenbucket.NewTokenBucketRateLimiter(1*time.Second, 5, 5)
+ _ = grpc.NewServer(
+ grpc_middleware.WithUnaryServerChain(
+ grpc_ratelimit.UnaryServerInterceptor(
+ grpc_ratelimit.WithLimiter(unaryRateLimiter),
+ grpc_ratelimit.WithMaxWaitDuration(10*time.Second),
+ ),
+ ),
+ grpc_middleware.WithStreamServerChain(
+ grpc_ratelimit.StreamServerInterceptor(
+ grpc_ratelimit.WithLimiter(streamRateLimiter),
+ grpc_ratelimit.WithMaxWaitDuration(5*time.Second),
+ ),
+ ),
+ )
+}

ratelimit/options.go

@@ -0,0 +1,27 @@
+// Copyright 2018 Zheng Dayu. All Rights Reserved.
+// See LICENSE for licensing terms.
+
+package grpc_ratelimit
+
+import "time"
+
+const infinityDuration time.Duration = 0x7fffffffffffffff
+
+type Option func(*rateLimiter)
+
+// WithLimiter customizes your limiter in the middleware
+func WithLimiter(l Limiter) Option {
+ return func(r *rateLimiter) {
+ r.limiter = l
+ if r.maxWaitDuration == 0 {
+ r.maxWaitDuration = infinityDuration
+ }
+ }
+}
+
+// WithMaxWaitDuration customizes maxWaitDuration in limiter's WaitMaxDuration action.
+func WithMaxWaitDuration(maxWaitDuration time.Duration) Option {
+ return func(r *rateLimiter) {
+ r.maxWaitDuration = maxWaitDuration
+ }
+}

ratelimit/ratelimit.go

@@ -0,0 +1,66 @@
+// Copyright 2018 Zheng Dayu. All Rights Reserved.
+// See LICENSE for licensing terms.
+
+package grpc_ratelimit
+
+import (
+ "time"
+
+ "golang.org/x/net/context"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/codes"
+ "google.golang.org/grpc/status"
+)
+
+type Limiter interface {
+ WaitMaxDuration(time.Duration) bool
+}
+
+type rateLimiter struct {
+ limiter Limiter
+ maxWaitDuration time.Duration
+}
+
+func (r *rateLimiter) Wait() bool {
+ return r.limiter.WaitMaxDuration(r.maxWaitDuration)
+}
+
+type emptyLimiter struct{}
+
+func (e *emptyLimiter) WaitMaxDuration(time.Duration) bool {
+ return true
+}
+
+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.Wait() {
+ return handler(ctx, req)
+ }
+ return nil, status.Errorf(codes.ResourceExhausted, "%s is rejected by grpc_ratelimit middleare, please retry later.", info.FullMethod)
+ }
+}
+
+// StreamServerInterceptor returns a new stream server interceptors that performs request rate limit.
+func StreamServerInterceptor(opts ...Option) grpc.StreamServerInterceptor {
+ ratelimiter := emptyRatelimiter()
+ for _, opt := range opts {
+ opt(ratelimiter)
+ }
+ return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+ if ratelimiter.Wait() {
+ return handler(srv, stream)
+ }
+ return status.Errorf(codes.ResourceExhausted, "%s is rejected by grpc_ratelimit middleare, please retry later.", info.FullMethod)
+ }
+}

ratelimit/ratelimit_test.go

@@ -0,0 +1,86 @@
+package grpc_ratelimit
+
+import (
+ "context"
+ "errors"
+ "testing"
+ "time"
+
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/metadata"
+
+ "github.com/grpc-ecosystem/go-grpc-middleware/ratelimit/tokenbucket"
+ "github.com/stretchr/testify/assert"
+)
+
+const errMsgFake = "fake error"
+
+func TestEmptyUnaryServerInterceptor(t *testing.T) {
+ interceptor := UnaryServerInterceptor()
+ handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+ return nil, errors.New(errMsgFake)
+ }
+ var ctx context.Context
+ var req interface{}
+ var info *grpc.UnaryServerInfo
+ req2, err := interceptor(ctx, req, info, handler)
+ assert.Nil(t, req2)
+ assert.EqualError(t, err, errMsgFake)
+}
+
+func TestRateLimitUnaryServerInterceptor(t *testing.T) {
+ unaryRateLimiter := tokenbucket.NewTokenBucketRateLimiter(1*time.Second, 1, 1)
+ interceptor := UnaryServerInterceptor(
+ WithLimiter(unaryRateLimiter),
+ WithMaxWaitDuration(1*time.Millisecond),
+ )
+ handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+ return nil, errors.New(errMsgFake)
+ }
+ var ctx context.Context
+ var req interface{}
+ var info *grpc.UnaryServerInfo
+ req2, err := interceptor(ctx, req, info, handler)
+ assert.Nil(t, req2)
+ assert.EqualError(t, err, errMsgFake)
+}
+
+func TestRateLimitStreamServerInterceptor(t *testing.T) {
+ unaryRateLimiter := tokenbucket.NewTokenBucketRateLimiter(1*time.Second, 1, 1)
+ interceptor := StreamServerInterceptor(
+ WithLimiter(unaryRateLimiter),
+ WithMaxWaitDuration(1*time.Millisecond),
+ )
+ handler := func(srv interface{}, stream grpc.ServerStream) error {
+ return errors.New(errMsgFake)
+ }
+ var srv interface{}
+ var ss *mockServerStream
+ var info *grpc.StreamServerInfo
+ err := interceptor(srv, ss, info, handler)
+ assert.EqualError(t, err, errMsgFake)
+}
+
+type mockServerStream struct{}
+
+func (mss *mockServerStream) SetHeader(metadata.MD) error {
+ return nil
+}
+
+func (mss *mockServerStream) SendHeader(metadata.MD) error {
+ return nil
+}
+
+func (mss *mockServerStream) SetTrailer(metadata.MD) {}
+
+func (mss *mockServerStream) Context() context.Context {
+ return context.Background()
+}
+
+func (mss *mockServerStream) SendMsg(m interface{}) error {
+ return nil
+}
+
+func (mss *mockServerStream) RecvMsg(m interface{}) error {
+ return nil
+}

ratelimit/tokenbucket/tokenbucket.go

@@ -0,0 +1,26 @@
+// 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
+}
+
+// NewTokenBucketRateLimiter creates a tokenBucketLimiter.
+func NewTokenBucketRateLimiter(fillInterval time.Duration, capacity, quantum int64) *tokenBucketLimiter {
+ return &tokenBucketLimiter{
+ limiter: ratelimit.NewBucketWithQuantum(fillInterval, capacity, quantum),
+ }
+}
+
+// WaitMaxDuration
+func (b *tokenBucketLimiter) WaitMaxDuration(maxWaitDuration time.Duration) bool {
+ return b.limiter.WaitMaxDuration(1, maxWaitDuration)
+}

ratelimit/tokenbucket/tokenbucket_test.go

@@ -0,0 +1,14 @@
+package tokenbucket
+
+import (
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+)
+
+func TestNewTokenBucketRateLimiter(t *testing.T) {
+ l := NewTokenBucketRateLimiter(1*time.Millisecond, 1, 1)
+ ok := l.WaitMaxDuration(1 * time.Millisecond)
+ assert.True(t, ok)
+}