ambassador: fix rate limit CRDs being duplicated (#168)

kubeshop · Sep 17, 2021 · 4c1ef48 · 4c1ef48
1 parent 5f6d268
commit 4c1ef48
Show file tree

Hide file tree

Showing 4 changed files with 152 additions and 34 deletions.
diff --git a/generators/ambassador/ambassador.go b/generators/ambassador/ambassador.go
@@ -109,7 +109,7 @@ func (g *Generator) Generate(opts *options.Options, spec *openapi3.T) (string, e
 	}
 
 	var mappings []mappingTemplateData
-	var rateLimits []rateLimitTemplateData
+	rateLimits := make(map[string]*rateLimitTemplateData)
 
 	serviceURL := g.getServiceURL(opts)
 
@@ -215,17 +215,36 @@ func (g *Generator) Generate(opts *options.Options, spec *openapi3.T) (string, e
 						}
 					}
 
-					if rateLimitOpts.Group == "" {
-						rateLimitOpts.Group = "default"
+					if rateLimitOpts.Group != "" {
+						// rate limit uses group, check that it wasn't already configured
+						if rl, ok := rateLimits[rateLimitOpts.Group]; ok {
+							// rate limit already configured within this group, replace limits if new ones are lower
+							if burstFactor < rl.BurstFactor {
+								rl.BurstFactor = burstFactor
+							}
+
+							if rps < rl.Rate {
+								rl.Rate = rps
+							}
+						} else {
+							rateLimits[rateLimitOpts.Group] = &rateLimitTemplateData{
+								Name:        opts.Service.Name + "-" + rateLimitOpts.Group,
+								Operation:   mappingName,
+								Rate:        rps,
+								BurstFactor: burstFactor,
+								Group:       rateLimitOpts.Group,
+							}
+						}
+					} else {
+						// rate limit on this operation does not use grouping
+						rateLimits[mappingName] = &rateLimitTemplateData{
+							Name:        opts.Service.Name + "-" + mappingName,
+							Operation:   mappingName,
+							Rate:        rps,
+							BurstFactor: burstFactor,
+						}
 					}
 
-					rateLimits = append(rateLimits, rateLimitTemplateData{
-						Operation:   mappingName,
-						Rate:        rps,
-						BurstFactor: burstFactor,
-						Group:       rateLimitOpts.Group,
-					})
-
 					op.LabelsEnabled = true
 					op.RateLimitGroup = rateLimitOpts.Group
 				}
@@ -301,12 +320,13 @@ func (g *Generator) Generate(opts *options.Options, spec *openapi3.T) (string, e
 				}
 			}
 
-			rateLimits = append(rateLimits, rateLimitTemplateData{
+			rateLimits["default"] = &rateLimitTemplateData{
+				Name:        "default",
 				Operation:   opts.Service.Name,
 				Rate:        rps,
 				BurstFactor: burstFactor,
 				Group:       opts.RateLimits.Group,
-			})
+			}
 		}
 
 		mappings = append(mappings, op)
@@ -320,8 +340,14 @@ func (g *Generator) Generate(opts *options.Options, spec *openapi3.T) (string, e
 		return mappings[i].MappingName < mappings[j].MappingName
 	})
 
-	sort.Slice(rateLimits, func(i, j int) bool {
-		return rateLimits[i].Operation < rateLimits[j].Operation
+	// flat out rate limits from map to array to sort them
+	var rateLimitsArray []rateLimitTemplateData
+	for _, rl := range rateLimits {
+		rateLimitsArray = append(rateLimitsArray, *rl)
+	}
+
+	sort.Slice(rateLimitsArray, func(i, j int) bool {
+		return rateLimitsArray[i].Operation < rateLimitsArray[j].Operation
 	})
 
 	var buf bytes.Buffer

diff --git a/generators/ambassador/ambassador_test.go b/generators/ambassador/ambassador_test.go
@@ -1309,21 +1309,18 @@ spec:
     ambassador:
 	  - group:
 		  - kusk-group-default
-      - operation:
-          - kusk-operation-petstore
       - request:
           - remote-address
 ---
 apiVersion: getambassador.io/v2
 kind: RateLimit
 metadata:
-  name: petstore
+  name: default
 spec:
   domain: ambassador
   limits:
     - pattern:
       - "generic_key": "kusk-group-default"
-        "generic_key": "kusk-operation-petstore"
         "remote-address": "*"
       rate: 100
       burstFactor: 2
@@ -1391,8 +1388,6 @@ spec:
   rewrite: ""
   labels:
     ambassador:
-	  - group:
-		  - kusk-group-default
       - operation:
           - kusk-operation-petstore-updatepet
       - request:
@@ -1411,8 +1406,6 @@ spec:
   rewrite: ""
   labels:
     ambassador:
-	  - group:
-		  - kusk-group-default
       - operation:
           - kusk-operation-petstore-uploadfile
       - request:
@@ -1421,13 +1414,12 @@ spec:
 apiVersion: getambassador.io/v2
 kind: RateLimit
 metadata:
-  name: petstore-updatepet
+  name: petstore-petstore-updatepet
 spec:
   domain: ambassador
   limits:
     - pattern:
-      - "generic_key": "kusk-group-default"
-        "generic_key": "kusk-operation-petstore-updatepet"
+      - "generic_key": "kusk-operation-petstore-updatepet"
         "remote-address": "*"
       rate: 20
       burstFactor: 2
@@ -1436,17 +1428,117 @@ spec:
 apiVersion: getambassador.io/v2
 kind: RateLimit
 metadata:
-  name: petstore-uploadfile
+  name: petstore-petstore-uploadfile
 spec:
   domain: ambassador
   limits:
     - pattern:
-      - "generic_key": "kusk-group-default"
-        "generic_key": "kusk-operation-petstore-uploadfile"
+      - "generic_key": "kusk-operation-petstore-uploadfile"
         "remote-address": "*"
       rate: 40
       burstFactor: 2
       unit: second
+`,
+		},
+		{
+			name: "rate-limit-group",
+			options: options.Options{
+				Namespace: "default",
+				Service: options.ServiceOptions{
+					Namespace: "default",
+					Name:      "petstore",
+				},
+				RateLimits: options.RateLimitOptions{
+					Group: "xyz",
+					RPS:   40,
+					Burst: 80,
+				},
+				PathSubOptions: map[string]options.SubOptions{
+					"/pet": {
+						RateLimits: options.RateLimitOptions{
+							Group: "xyz",
+							RPS:   20,
+							Burst: 40,
+						},
+					},
+				},
+			},
+			spec: `
+openapi: 3.0.2
+info:
+  title: Swagger Petstore - OpenAPI 3.0
+  version: 1.0.5
+paths:
+  "/pet":
+    put:
+      operationId: updatePet
+      responses:
+        '200':
+          description: Successful operation
+  "/pet/{petId}/uploadImage":
+    post:
+      operationId: uploadFile
+      parameters:
+        - name: petId
+          in: path
+          description: ID of pet to update
+          required: true
+          schema:
+            type: integer
+            format: int64
+      responses:
+        '200':
+          description: Successful operation`,
+			res: `
+---
+apiVersion: getambassador.io/v2
+kind: Mapping
+metadata:
+  name: petstore-updatepet
+  namespace: default
+spec:
+  prefix: "/pet"
+  method: PUT
+  service: petstore.default:80
+  rewrite: ""
+  labels:
+    ambassador:
+	  - group:
+		  - kusk-group-xyz
+      - request:
+          - remote-address
+---
+apiVersion: getambassador.io/v2
+kind: Mapping
+metadata:
+  name: petstore-uploadfile
+  namespace: default
+spec:
+  prefix: "/pet/([a-zA-Z0-9]*)/uploadImage"
+  prefix_regex: true
+  method: POST
+  service: petstore.default:80
+  rewrite: ""
+  labels:
+    ambassador:
+	  - group:
+		  - kusk-group-xyz
+      - request:
+          - remote-address
+---
+apiVersion: getambassador.io/v2
+kind: RateLimit
+metadata:
+  name: petstore-xyz
+spec:
+  domain: ambassador
+  limits:
+    - pattern:
+      - "generic_key": "kusk-group-xyz"
+        "remote-address": "*"
+      rate: 20
+      burstFactor: 2
+      unit: second
 `,
 		},
 		{

diff --git a/generators/ambassador/mapping_template.go b/generators/ambassador/mapping_template.go
@@ -83,11 +83,11 @@ spec:
 
   {{if .LabelsEnabled}}
   labels:
-    ambassador:
+    ambassador:{{if .RateLimitGroup}}
 	  - group:
-		  - kusk-group-{{.RateLimitGroup}}
+		  - kusk-group-{{.RateLimitGroup}}{{else}}
       - operation:
-          - kusk-operation-{{.MappingName}}
+          - kusk-operation-{{.MappingName}}{{end}}
       - request:
           - remote-address
   {{end}}

diff --git a/generators/ambassador/rate_limit_template.go b/generators/ambassador/rate_limit_template.go
@@ -2,6 +2,7 @@ package ambassador
 
 type rateLimitTemplateData struct {
 	Group       string
+	Name        string
 	Operation   string
 	Rate        uint32
 	BurstFactor uint32
@@ -12,13 +13,12 @@ var rateLimitTemplateRaw = `{{range .}}
 apiVersion: getambassador.io/v2
 kind: RateLimit
 metadata:
-  name: {{.Operation}}
+  name: {{.Name}}
 spec:
   domain: ambassador
   limits:
     - pattern:
-      - "generic_key": "kusk-group-{{.Group}}"
-        "generic_key": "kusk-operation-{{.Operation}}"
+      - {{if .Group}}"generic_key": "kusk-group-{{.Group}}"{{else}}"generic_key": "kusk-operation-{{.Operation}}"{{end}}
         "remote-address": "*"
       rate: {{.Rate}}
       {{if .BurstFactor}}