RHINENG-2426: add leaky bucket rate limiter

RedHatInsights · Oct 19, 2023 · a71a085 · a71a085
1 parent 2d4769e
commit a71a085
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 0 deletions.
diff --git a/base/utils/config.go b/base/utils/config.go
@@ -40,6 +40,7 @@ type Config struct {
 	MaxRequestBodySize int64
 	MaxHeaderCount     int
 	MaxGinConnections  int
+	Ratelimit          int
 
 	// kafka
 	KafkaAddress           string
@@ -174,6 +175,7 @@ func initAPIromClowder() {
 	Cfg.MaxRequestBodySize = GetInt64EnvOrDefault("MAX_REQUEST_BODY_SIZE", 1*1024*1024)
 	Cfg.MaxHeaderCount = GetIntEnvOrDefault("MAX_HEADER_COUNT", 50)
 	Cfg.MaxGinConnections = GetIntEnvOrDefault("MAX_GIN_CONNECTIONS", 50)
+	Cfg.Ratelimit = GetIntEnvOrDefault("RATELIMIT", 100)
 }
 
 func initKafkaFromClowder() {

diff --git a/deploy/clowdapp.yaml b/deploy/clowdapp.yaml
@@ -126,6 +126,7 @@ objects:
         - {name: MAX_REQUEST_BODY_SIZE, value: '${MAX_REQUEST_BODY_SIZE}'}
         - {name: MAX_HEADER_COUNT, value: '${MAX_HEADER_COUNT}'}
         - {name: MAX_GIN_CONNECTIONS, value: '${MAX_GIN_CONNECTIONS}'}
+        - {name: RATELIMIT, value: '${RATELIMIT}'}
 
         resources:
           limits: {cpu: '${RES_LIMIT_CPU_MANAGER}', memory: '${RES_LIMIT_MEM_MANAGER}'}
@@ -583,6 +584,7 @@ parameters:
 - {name: MAX_REQUEST_BODY_SIZE, value: '1048576'} # limit request body size, in bytes (default 1MB)
 - {name: MAX_HEADER_COUNT, value: '50'} # limit number of request headers
 - {name: MAX_GIN_CONNECTIONS, value: '50'}
+- {name: RATELIMIT, value: '100'} # requests per second for leaky bucket rate limiter
 
 # Listener
 - {name: REPLICAS_LISTENER, value: '1'}

diff --git a/go.mod b/go.mod
@@ -28,6 +28,7 @@ require (
 	github.com/swaggo/gin-swagger v1.6.0
 	github.com/zsais/go-gin-prometheus v0.1.0
 	go.uber.org/automaxprocs v1.5.2
+	go.uber.org/ratelimit v0.3.0
 	golang.org/x/net v0.17.0
 	gorm.io/driver/postgres v1.5.2
 	gorm.io/gorm v1.25.1
@@ -36,6 +37,7 @@ require (
 
 require (
 	github.com/KyleBanks/depth v1.2.1 // indirect
+	github.com/benbjohnson/clock v1.3.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bytedance/sonic v1.9.2 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect

diff --git a/go.sum b/go.sum
@@ -7,6 +7,8 @@ github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migc
 github.com/aws/aws-sdk-go v1.38.51/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zKMmprdro=
 github.com/aws/aws-sdk-go v1.44.280 h1:UYl/yxhDxP8naok6ftWyQ9/9ZzNwjC9dvEs/j8BkGhw=
 github.com/aws/aws-sdk-go v1.44.280/go.mod h1:aVsgQcEevwlmQ7qHE9I3h+dtQgpqhFB+i8Phjh7fkwI=
+github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
+github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
@@ -276,6 +278,8 @@ go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0
 go.uber.org/automaxprocs v1.5.2 h1:2LxUOGiR3O6tw8ui5sZa2LAaHnsviZdVOUZw4fvbnME=
 go.uber.org/automaxprocs v1.5.2/go.mod h1:eRbA25aqJrxAbsLO0xy5jVwPt7FQnRgjW+efnwa1WM0=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
+go.uber.org/ratelimit v0.3.0 h1:IdZd9wqvFXnvLvSEBo0KPcGfkoBGNkpTHlrE3Rcjkjw=
+go.uber.org/ratelimit v0.3.0/go.mod h1:So5LG7CV1zWpY1sHe+DXTJqQvOx+FFPFaAs2SnoyBaI=
 go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
 golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
 golang.org/x/arch v0.4.0 h1:A8WCeEWhLwPBKNbFi5Wv5UTCBx5zzubnXDlMOFAzFMc=

diff --git a/manager/manager.go b/manager/manager.go
@@ -45,6 +45,7 @@ func RunManager() {
 	app.Use(gin.Recovery())
 	middlewares.Prometheus().Use(app)
 	app.Use(middlewares.MaxConnections(utils.Cfg.MaxGinConnections))
+	app.Use(middlewares.Ratelimit(utils.Cfg.Ratelimit))
 	app.Use(middlewares.RequestResponseLogger())
 	app.Use(gzip.Gzip(gzip.DefaultCompression))
 	endpointsConfig := getEndpointsConfig()

diff --git a/manager/middlewares/limits.go b/manager/middlewares/limits.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
+	"go.uber.org/ratelimit"
 )
 
 func LimitRequestBodySize(size int64) gin.HandlerFunc {
@@ -32,3 +33,11 @@ func MaxConnections(max int) gin.HandlerFunc {
 		c.Next()
 	}
 }
+
+func Ratelimit(max int) gin.HandlerFunc {
+	rl := ratelimit.New(max)
+	return func(c *gin.Context) {
+		rl.Take()
+		c.Next()
+	}
+}