add rate limit config

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 	github.com/hashicorp/consul/api v1.8.1
+	github.com/juju/ratelimit v1.0.1
 	github.com/mbobakov/grpc-consul-resolver v1.4.4
 	github.com/opentracing/opentracing-go v1.2.0
 	github.com/prometheus/client_golang v1.11.0

go.sum

@@ -155,6 +155,8 @@ github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCV
 github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
 github.com/json-iterator/go v1.1.10/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
 github.com/json-iterator/go v1.1.11/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
+github.com/juju/ratelimit v1.0.1 h1:+7AIFJVQ0EQgq/K9+0Krm7m530Du7tIz0METWzN0RgY=
+github.com/juju/ratelimit v1.0.1/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk=
 github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
 github.com/julienschmidt/httprouter v1.3.0/go.mod h1:JR6WtHb+2LUe8TCKY3cZOxFyyO8IZAc4RVcycCCAKdM=
 github.com/jung-kurt/gofpdf v1.0.3-0.20190309125859-24315acbbda5/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes=

rpc/config.go

@@ -26,6 +26,7 @@ const (
 type Config struct {
 	Pprof        pprofConfig `toml:"pprof"`
 	Server       serverConfig
+	RateLimit    rateLimitConfig `toml:"rate_limit"`
 	Consul       consulConfig
 	Metrics      metricsConfig
 	Trace        traceConfig
@@ -41,6 +42,13 @@ type serverConfig struct {
 	HttpPort    int    `toml:"http_port"`
 }
 
+type rateLimitConfig struct {
+	Enabled      bool
+	Type         string `toml:"type" default:"always_pass"`
+	FillInterval int    `toml:"fill_interval" default:"300"`
+	Capacity     int64  `toml:"capacity" default:"3000"`
+}
+
 type pprofConfig struct {
 	Port int `toml:"port"`
 }

rpc/ratelimit.go

@@ -1,9 +1,55 @@
 package rpc
 
+import (
+	"time"
+
+	grpcRateLimit "github.com/grpc-ecosystem/go-grpc-middleware/ratelimit"
+	"github.com/juju/ratelimit"
+)
+
+const (
+	LimiterAlwaysPass = "always_pass"
+	LimiterNoBlock    = "no_block"
+)
+
+var limiterMap = map[string]grpcRateLimit.Limiter{
+	LimiterAlwaysPass: &alwaysPassLimiter{},
+	LimiterNoBlock:    &rateLimitNoBlock{},
+}
+
+func name2Limiter(typeName string) grpcRateLimit.Limiter {
+	return limiterMap[typeName]
+}
+
 // alwaysPassLimiter is an example limiter which implements Limiter interface.
 // It does not limit any request because Limit function always returns false.
 type alwaysPassLimiter struct{}
 
 func (*alwaysPassLimiter) Limit() bool {
 	return false
 }
+
+var tb *ratelimit.Bucket
+
+func initRateLimit(cfg *Config) grpcRateLimit.Limiter {
+	if cfg.RateLimit.Type != LimiterAlwaysPass {
+		tb = ratelimit.NewBucket(time.Duration(cfg.RateLimit.FillInterval), cfg.RateLimit.Capacity)
+	}
+	limiter := name2Limiter(cfg.RateLimit.Type)
+	return limiter
+}
+
+// rateLimitNoBlock 如果桶中没有token，不block，直接返回
+type rateLimitNoBlock struct{}
+
+func (*rateLimitNoBlock) Limit() bool {
+	if tb == nil {
+		return false
+	}
+	count := tb.TakeAvailable(1)
+	if count == 0 {
+		return true
+	} else {
+		return false
+	}
+}

rpc/server.go

@@ -19,10 +19,10 @@ import (
 )
 
 type Server struct {
-	cfg         *Config
-	server      *grpc.Server
-	Log         Logger
-	Err         error
+	cfg    *Config
+	server *grpc.Server
+	Log    Logger
+	Err    error
 }
 
 func initGrpcServer(cfg *Config) *grpc.Server {
@@ -52,9 +52,11 @@ func makeMiddlewareInterceptor(cfg *Config) []grpc.ServerOption {
 	}
 
 	// rate limit
-	limiter := &alwaysPassLimiter{}
-	siList = append(siList, ratelimit.StreamServerInterceptor(limiter))
-	uiList = append(uiList, ratelimit.UnaryServerInterceptor(limiter))
+	if cfg.RateLimit.Enabled {
+		limiter := initRateLimit(cfg)
+		siList = append(siList, ratelimit.StreamServerInterceptor(limiter))
+		uiList = append(uiList, ratelimit.UnaryServerInterceptor(limiter))
+	}
 
 	// panic recovery
 	opts := []grpc_recovery.Option{