rate limit: add group option (#153)

kubeshop · Sep 2, 2021 · 5b96f6f · 5b96f6f
1 parent a6f701a
commit 5b96f6f
Show file tree

Hide file tree

Showing 7 changed files with 52 additions and 4 deletions.
diff --git a/docs/ambassador.md b/docs/ambassador.md
@@ -42,6 +42,7 @@ To override settings on the path or HTTP method level, you are required to use t
 | Host                    | --host                     | host                      | The value to set the host field to in the Mapping resource                                                         | ✅                             |
 | Rate limit (RPS)        | --rate_limits.rps          | rate_limits.rps           | Request per second rate limit                                                                                      | ✅                             |
 | Rate limit (burst)      | --rate_limits.burst        | rate_limits.burst         | Rate limit burst                                                                                                   | ✅                             |
+| Rate limit group        | N/A                        | rate_limits.group         | Rate limit endpoint group                                                                                          |                               |
 | Request Timeout         | --timeouts.request_timeout | timeouts.request_timeout  | Total request timeout (seconds)                                                                                    | ✅                             |
 | Idle Timeout            | --timeouts.idle_timeout    | timeouts.idle_timeout     | Idle connection timeout (seconds)                                                                                  | ✅                             |
 | CORS Origins            | N/A                        | cors.origins              | Array of origins                                                                                                   | ✅                             |

diff --git a/generators/ambassador/ambassador.go b/generators/ambassador/ambassador.go
@@ -222,13 +222,19 @@ func (g *Generator) Generate(opts *options.Options, spec *openapi3.T) (string, e
 						}
 					}
 
+					if rateLimitOpts.Group == "" {
+						rateLimitOpts.Group = "default"
+					}
+
 					rateLimits = append(rateLimits, rateLimitTemplateData{
 						Operation:   mappingName,
 						Rate:        rps,
 						BurstFactor: burstFactor,
+						Group:       rateLimitOpts.Group,
 					})
 
 					op.LabelsEnabled = true
+					op.RateLimitGroup = rateLimitOpts.Group
 				}
 
 				// take global timeout options
@@ -281,6 +287,12 @@ func (g *Generator) Generate(opts *options.Options, spec *openapi3.T) (string, e
 		if !reflect.DeepEqual(options.RateLimitOptions{}, opts.RateLimits) {
 			op.LabelsEnabled = true
 
+			if opts.RateLimits.Group == "" {
+				opts.RateLimits.Group = "default"
+			}
+
+			op.RateLimitGroup = opts.RateLimits.Group
+
 			rps := opts.RateLimits.RPS
 
 			var burstFactor uint32
@@ -300,6 +312,7 @@ func (g *Generator) Generate(opts *options.Options, spec *openapi3.T) (string, e
 				Operation:   opts.Service.Name,
 				Rate:        rps,
 				BurstFactor: burstFactor,
+				Group:       opts.RateLimits.Group,
 			})
 		}
 

diff --git a/generators/ambassador/ambassador_test.go b/generators/ambassador/ambassador_test.go
@@ -1320,6 +1320,8 @@ spec:
   rewrite: ""
   labels:
     ambassador:
+	  - group:
+		  - kusk-group-default
       - operation:
           - kusk-operation-petstore
       - request:
@@ -1333,7 +1335,8 @@ spec:
   domain: ambassador
   limits:
     - pattern:
-      - "generic_key": "kusk-operation-petstore"
+      - "generic_key": "kusk-group-default"
+        "generic_key": "kusk-operation-petstore"
         "remote-address": "*"
       rate: 100
       burstFactor: 2
@@ -1401,6 +1404,8 @@ spec:
   rewrite: ""
   labels:
     ambassador:
+	  - group:
+		  - kusk-group-default
       - operation:
           - kusk-operation-petstore-updatepet
       - request:
@@ -1419,6 +1424,8 @@ spec:
   rewrite: ""
   labels:
     ambassador:
+	  - group:
+		  - kusk-group-default
       - operation:
           - kusk-operation-petstore-uploadfile
       - request:
@@ -1432,7 +1439,8 @@ spec:
   domain: ambassador
   limits:
     - pattern:
-      - "generic_key": "kusk-operation-petstore-updatepet"
+      - "generic_key": "kusk-group-default"
+        "generic_key": "kusk-operation-petstore-updatepet"
         "remote-address": "*"
       rate: 20
       burstFactor: 2
@@ -1446,7 +1454,8 @@ spec:
   domain: ambassador
   limits:
     - pattern:
-      - "generic_key": "kusk-operation-petstore-uploadfile"
+      - "generic_key": "kusk-group-default"
+        "generic_key": "kusk-operation-petstore-uploadfile"
         "remote-address": "*"
       rate: 40
       burstFactor: 2

diff --git a/generators/ambassador/mapping_template.go b/generators/ambassador/mapping_template.go
@@ -35,6 +35,8 @@ type mappingTemplateData struct {
 
 	RequestTimeout uint32
 	IdleTimeout    uint32
+
+	RateLimitGroup string
 }
 
 var mappingTemplateRaw = `{{range .}}
@@ -82,6 +84,8 @@ spec:
   {{if .LabelsEnabled}}
   labels:
     ambassador:
+	  - group:
+		  - kusk-group-{{.RateLimitGroup}}
       - operation:
           - kusk-operation-{{.MappingName}}
       - request:

diff --git a/generators/ambassador/rate_limit_template.go b/generators/ambassador/rate_limit_template.go
@@ -1,6 +1,7 @@
 package ambassador
 
 type rateLimitTemplateData struct {
+	Group       string
 	Operation   string
 	Rate        uint32
 	BurstFactor uint32
@@ -16,7 +17,8 @@ spec:
   domain: ambassador
   limits:
     - pattern:
-      - "generic_key": "kusk-operation-{{.Operation}}"
+      - "generic_key": "kusk-group-{{.Group}}"
+        "generic_key": "kusk-operation-{{.Operation}}"
         "remote-address": "*"
       rate: {{.Rate}}
       {{if .BurstFactor}}

diff --git a/generators/nginx_ingress/nginx_ingress.go b/generators/nginx_ingress/nginx_ingress.go
@@ -314,6 +314,22 @@ func (g *Generator) shouldSplit(opts *options.Options, spec *openapi3.T) bool {
 	}
 
 	rateLimitWarned := false
+	groupUnsupportedWarned := false
+
+	warnGroupUnsupported := func(ro options.RateLimitOptions) {
+		if groupUnsupportedWarned {
+			return
+		}
+
+		if ro.Group != "" {
+			log.New(os.Stderr, "WARN", log.Lmsgprefix).
+				Printf("nginx-ingress does not support rate limit groups. These will be ignored")
+
+			groupUnsupportedWarned = true
+		}
+	}
+
+	warnGroupUnsupported(opts.RateLimits)
 
 	for path, pathItem := range spec.Paths {
 		if pathSubOptions, ok := opts.PathSubOptions[path]; ok {
@@ -332,6 +348,8 @@ func (g *Generator) shouldSplit(opts *options.Options, spec *openapi3.T) bool {
 			if !reflect.DeepEqual(options.RateLimitOptions{}, pathSubOptions.RateLimits) &&
 				!reflect.DeepEqual(opts.RateLimits, pathSubOptions.RateLimits) {
 
+				warnGroupUnsupported(pathSubOptions.RateLimits)
+
 				if !rateLimitWarned {
 					log.New(os.Stderr, "WARN", log.Lmsgprefix).
 						Printf("Setting a rate limit option on the path level would cause a separate rate limit applied for each path")

diff --git a/options/ratelimit.go b/options/ratelimit.go
@@ -3,6 +3,7 @@ package options
 type RateLimitOptions struct {
 	RPS   uint32 `json:"rps,omitempty" yaml:"rps,omitempty"`
 	Burst uint32 `json:"burst,omitempty" yaml:"burst,omitempty"`
+	Group string `json:"group,omitempty" yaml:"group,omitempty"`
 }
 
 func (o *RateLimitOptions) Validate() error {