Add pgsql scenario (#3125)

- **New Features**
- Introduced a load testing script for PostgreSQL database concurrency.
- Added a new Dockerfile for building a Go application with health
check.
- Implemented PostgreSQL database connection and handling in the Go SDK
example.

- **Enhancements**
  - Search functionality now supports custom HTTP headers.
- Updated Kubernetes configurations with namespace, labels, and
environment variables.

- **Bug Fixes**
- Fixed the initialization of `inflightRequests` to prevent logic errors
in request handling.

- **Refactor**
- Replaced `wget` with `curl` for HTTP requests to improve reliability
and performance.
  - Modified `flow.go` to better manage inflight requests.

---------

Co-authored-by: Hardik Shingala <[email protected]>

blueprints/concurrency-limiting/base/config.libsonnet

@@ -13,7 +13,6 @@ commonConfig {
       max_concurrency: '__REQUIRED_FIELD__',
       selectors: commonConfig.selectors_defaults,
       parameters: {
-        limit_by_label_key: 'limit_by_label_key',
         max_inflight_duration: '__REQUIRED_FIELD__',
       },
       request_parameters: {},

blueprints/concurrency-limiting/base/gen/definitions.json

@@ -67,7 +67,6 @@
             "parameters": {
               "description": "Parameters.",
               "default": {
-                "limit_by_label_key": "limit_by_label_key",
                 "max_inflight_duration": "__REQUIRED_FIELD__"
               },
               "type": "object",

blueprints/concurrency-limiting/base/gen/values.yaml

@@ -29,7 +29,6 @@ policy:
     # Type: aperture.spec.v1.ConcurrencyLimiterParameters
     # Required: True
     parameters:
-      limit_by_label_key: "limit_by_label_key"
       max_inflight_duration: __REQUIRED_FIELD__
     # Request Parameters.
     # Type: aperture.spec.v1.ConcurrencyLimiterRequestParameters

blueprints/concurrency-scheduling/base/config.libsonnet

@@ -13,14 +13,9 @@ commonConfig {
       max_concurrency: '__REQUIRED_FIELD__',
       selectors: commonConfig.selectors_defaults,
       concurrency_limiter: {
-        limit_by_label_key: 'limit_by_label_key',
         max_inflight_duration: '__REQUIRED_FIELD__',
       },
-      scheduler: {
-        tokens_label_key: 'tokens',
-        priority_label_key: 'priority',
-        workload_label_key: 'workload',
-      },
+      scheduler: {},
       alerter: {
         alert_name: 'Too many inflight requests',
       },

blueprints/concurrency-scheduling/base/gen/definitions.json

@@ -61,7 +61,6 @@
             "concurrency_limiter": {
               "description": "Concurrency Limiter Parameters.",
               "default": {
-                "limit_by_label_key": "limit_by_label_key",
                 "max_inflight_duration": "__REQUIRED_FIELD__"
               },
               "type": "object",
@@ -75,11 +74,7 @@
             },
             "scheduler": {
               "description": "Scheduler configuration.",
-              "default": {
-                "priority_label_key": "priority",
-                "tokens_label_key": "tokens",
-                "workload_label_key": "workload"
-              },
+              "default": {},
               "type": "object",
               "$ref": "../../../gen/jsonschema/_definitions.json#/definitions/Scheduler"
             },

blueprints/concurrency-scheduling/base/gen/values.yaml

@@ -25,18 +25,14 @@ policy:
     # Type: aperture.spec.v1.ConcurrencyLimiterParameters
     # Required: True
     concurrency_limiter:
-      limit_by_label_key: "limit_by_label_key"
       max_inflight_duration: __REQUIRED_FIELD__
     # Max concurrency.
     # Type: float64
     # Required: True
     max_concurrency: __REQUIRED_FIELD__
     # Scheduler configuration.
     # Type: aperture.spec.v1.Scheduler
-    scheduler:
-      priority_label_key: "priority"
-      tokens_label_key: "tokens"
-      workload_label_key: "workload"
+    scheduler: {}
     # Flow selectors to match requests against.
     # Type: []aperture.spec.v1.Selector
     # Required: True

blueprints/quota-scheduling/base/config.libsonnet

@@ -15,14 +15,9 @@ commonConfig {
       fill_amount: '__REQUIRED_FIELD__',
       selectors: commonConfig.selectors_defaults,
       rate_limiter: {
-        limit_by_label_key: 'limit_key',
         interval: '__REQUIRED_FIELD__',
       },
-      scheduler: {
-        tokens_label_key: 'tokens',
-        priority_label_key: 'priority',
-        workload_label_key: 'workload',
-      },
+      scheduler: {},
       alerter: {
         alert_name: 'More than 90% of requests are being rate limited',
       },

blueprints/quota-scheduling/base/gen/definitions.json

@@ -78,19 +78,14 @@
             "rate_limiter": {
               "description": "Rate Limiter Parameters.",
               "default": {
-                "interval": "__REQUIRED_FIELD__",
-                "limit_by_label_key": "limit_key"
+                "interval": "__REQUIRED_FIELD__"
               },
               "type": "object",
               "$ref": "../../../gen/jsonschema/_definitions.json#/definitions/RateLimiterParameters"
             },
             "scheduler": {
               "description": "Scheduler configuration.",
-              "default": {
-                "priority_label_key": "priority",
-                "tokens_label_key": "tokens",
-                "workload_label_key": "workload"
-              },
+              "default": {},
               "type": "object",
               "$ref": "../../../gen/jsonschema/_definitions.json#/definitions/Scheduler"
             },

blueprints/quota-scheduling/base/gen/values.yaml

@@ -34,13 +34,9 @@ policy:
     # Required: True
     rate_limiter:
       interval: __REQUIRED_FIELD__
-      limit_by_label_key: "limit_key"
     # Scheduler configuration.
     # Type: aperture.spec.v1.Scheduler
-    scheduler:
-      priority_label_key: "priority"
-      tokens_label_key: "tokens"
-      workload_label_key: "workload"
+    scheduler: {}
     # Flow selectors to match requests against.
     # Type: []aperture.spec.v1.Selector
     # Required: True

blueprints/rate-limiting/base/config.libsonnet

@@ -15,7 +15,6 @@ commonConfig {
       fill_amount: '__REQUIRED_FIELD__',
       selectors: commonConfig.selectors_defaults,
       parameters: {
-        limit_by_label_key: 'limit_key',
         interval: '__REQUIRED_FIELD__',
       },
       request_parameters: {},

blueprints/rate-limiting/base/gen/definitions.json

@@ -78,8 +78,7 @@
             "parameters": {
               "description": "Parameters.",
               "default": {
-                "interval": "__REQUIRED_FIELD__",
-                "limit_by_label_key": "limit_key"
+                "interval": "__REQUIRED_FIELD__"
               },
               "type": "object",
               "$ref": "../../../gen/jsonschema/_definitions.json#/definitions/RateLimiterParameters"

blueprints/rate-limiting/base/gen/values.yaml

@@ -34,7 +34,6 @@ policy:
     # Required: True
     parameters:
       interval: __REQUIRED_FIELD__
-      limit_by_label_key: "limit_key"
     # Request Parameters.
     # Type: aperture.spec.v1.RateLimiterRequestParameters
     request_parameters: {}

docs/content/aperture-for-infra/guides/api-quota-management/assets/inter-service-rate-limiting/policy.yaml

@@ -26,9 +26,6 @@ spec:
               num_sync: 4
             limit_by_label_key: api_key
           scheduler:
-            priority_label_key: priority
-            tokens_label_key: tokens
-            workload_label_key: workload
             workloads:
             - label_matcher:
                 match_labels:

docs/content/guides/assets/managing-quotas/policy.yaml

@@ -24,7 +24,6 @@ spec:
             limit_by_label_key: user_id
           scheduler:
             priority_label_key: priority
-            tokens_label_key: tokens
             workload_label_key: workload
           selectors:
           - control_point: quota-scheduling-feature

docs/content/guides/assets/openai/policy.yaml

@@ -24,8 +24,6 @@ spec:
             limit_by_label_key: api_key
           scheduler:
             priority_label_key: priority
-            tokens_label_key: tokens
-            workload_label_key: workload
           selectors:
           - control_point: openai
             label_matcher:

docs/content/reference/aperture-cli/policies/assets/raw_values.yaml

@@ -34,7 +34,6 @@ policy:
     # Required: True
     parameters:
       interval: __REQUIRED_FIELD__
-      limit_by_label_key: "limit_key"
     # Request Parameters.
     # Type: aperture.spec.v1.RateLimiterRequestParameters
     request_parameters: {}

docs/content/reference/blueprints/concurrency-limiting/base.md

@@ -121,7 +121,7 @@ href={`https://github.com/fluxninja/aperture/tree/${aver}/blueprints/concurrency
     description='Parameters.'
     type='Object (aperture.spec.v1.ConcurrencyLimiterParameters)'
     reference='../../configuration/spec#concurrency-limiter-parameters'
-    value='{"limit_by_label_key": "limit_by_label_key", "max_inflight_duration": "__REQUIRED_FIELD__"}'
+    value='{"max_inflight_duration": "__REQUIRED_FIELD__"}'
 />
 
 <!-- vale on -->

docs/content/reference/blueprints/concurrency-scheduling/base.md

@@ -107,7 +107,7 @@ href={`https://github.com/fluxninja/aperture/tree/${aver}/blueprints/concurrency
     description='Concurrency Limiter Parameters.'
     type='Object (aperture.spec.v1.ConcurrencyLimiterParameters)'
     reference='../../configuration/spec#concurrency-limiter-parameters'
-    value='{"limit_by_label_key": "limit_by_label_key", "max_inflight_duration": "__REQUIRED_FIELD__"}'
+    value='{"max_inflight_duration": "__REQUIRED_FIELD__"}'
 />
 
 <!-- vale on -->
@@ -135,7 +135,7 @@ href={`https://github.com/fluxninja/aperture/tree/${aver}/blueprints/concurrency
     description='Scheduler configuration.'
     type='Object (aperture.spec.v1.Scheduler)'
     reference='../../configuration/spec#scheduler'
-    value='{"priority_label_key": "priority", "tokens_label_key": "tokens", "workload_label_key": "workload"}'
+    value='{}'
 />
 
 <!-- vale on -->

docs/content/reference/blueprints/quota-scheduling/base.md

@@ -135,7 +135,7 @@ href={`https://github.com/fluxninja/aperture/tree/${aver}/blueprints/quota-sched
     description='Rate Limiter Parameters.'
     type='Object (aperture.spec.v1.RateLimiterParameters)'
     reference='../../configuration/spec#rate-limiter-parameters'
-    value='{"interval": "__REQUIRED_FIELD__", "limit_by_label_key": "limit_key"}'
+    value='{"interval": "__REQUIRED_FIELD__"}'
 />
 
 <!-- vale on -->
@@ -149,7 +149,7 @@ href={`https://github.com/fluxninja/aperture/tree/${aver}/blueprints/quota-sched
     description='Scheduler configuration.'
     type='Object (aperture.spec.v1.Scheduler)'
     reference='../../configuration/spec#scheduler'
-    value='{"priority_label_key": "priority", "tokens_label_key": "tokens", "workload_label_key": "workload"}'
+    value='{}'
 />
 
 <!-- vale on -->

docs/content/reference/blueprints/rate-limiting/base.md

@@ -135,7 +135,7 @@ href={`https://github.com/fluxninja/aperture/tree/${aver}/blueprints/rate-limiti
     description='Parameters.'
     type='Object (aperture.spec.v1.RateLimiterParameters)'
     reference='../../configuration/spec#rate-limiter-parameters'
-    value='{"interval": "__REQUIRED_FIELD__", "limit_by_label_key": "limit_key"}'
+    value='{"interval": "__REQUIRED_FIELD__"}'
 />
 
 <!-- vale on -->

pkg/scheduler/wfq.go

@@ -528,11 +528,32 @@ func (pMetrics *preemptionMetrics) onQueueEntry(request *Request, qRequest *queu
 // Update metrics for preemption and delay
 // WARNING: Unsafe and should be called with scheduler lock.
 func (pMetrics *preemptionMetrics) onQueueExit(request *Request, qRequest *queuedRequest, allowed bool) {
+	initMetrics := func(labels prometheus.Labels) error {
+		var err error
+		_, err = pMetrics.workloadPreemptedTokensSummary.GetMetricWith(labels)
+		if err != nil {
+			return fmt.Errorf("%w: failed to get workload_preempted_tokens summary", err)
+		}
+		_, err = pMetrics.workloadDelayedTokensSummary.GetMetricWith(labels)
+		if err != nil {
+			return fmt.Errorf("%w: failed to get workload_delayed_tokens summary", err)
+		}
+		_, err = pMetrics.workloadOnTimeCounter.GetMetricWith(labels)
+		if err != nil {
+			return fmt.Errorf("%w: failed to get workload_on_time_total counter", err)
+		}
+		return nil
+	}
+
 	publishSummary := func(summary *prometheus.SummaryVec, value float64) {
 		if summary == nil {
 			return
 		}
 		metricsLabels := appendWorkloadLabel(pMetrics.metricsLabels, request.FairnessLabel)
+		err := initMetrics(metricsLabels)
+		if err != nil {
+			log.Error().Err(err).Msg("Failed to initialize metrics")
+		}
 		observer, err := summary.GetMetricWith(metricsLabels)
 		if err != nil {
 			log.Error().Err(err).Msg("Failed to get workload preempted tokens summary")
@@ -546,6 +567,10 @@ func (pMetrics *preemptionMetrics) onQueueExit(request *Request, qRequest *queue
 			return
 		}
 		metricsLabels := appendWorkloadLabel(pMetrics.metricsLabels, request.FairnessLabel)
+		err := initMetrics(metricsLabels)
+		if err != nil {
+			log.Error().Err(err).Msg("Failed to initialize metrics")
+		}
 		counter, err := counterVec.GetMetricWith(metricsLabels)
 		if err != nil {
 			log.Error().Err(err).Msg("Failed to get workload on time counter")

playground/Tiltfile

@@ -379,6 +379,9 @@ DEP_TREE = {
                 {
                     "workload": "aperture-go-example",
                     "resource_deps": ["agent", "controller", "cluster-bootstrap"],
+                    "extra_objects": [
+                        ("^aperture-go-example", "serviceaccount", APERTURE_GO_EXAMPLE_NS)
+                    ],
                 }
             ]
         },
@@ -1133,6 +1136,10 @@ def declare_resources(resources, dep_tree, inv_dep_tree, race_arg, cloud_extensi
             namespace = app_config.get("namespace", None)
             if tkenv:
                 base_dir = app_config.get("base_dir")
+                if not namespace and "playground/tanka/apps/aperture-go-example" in tkenv and APERTURE_GO_EXAMPLE_NS + ":namespace" not in managed_namespaces:
+                    namespace_create(APERTURE_GO_EXAMPLE_NS, labels=ns_labels)
+                    namespace = APERTURE_GO_EXAMPLE_NS
+                    managed_namespaces += [namespace + ":namespace"]
                 render_tanka(tkenv, env_vars, base_dir, namespace, cloud_extension, values)
             elif renderer == "yaml":
                 render_yaml(manifests_path, env_vars)

playground/scenarios/concurrency-scheduler/policies/concurrency-scheduler-cr.yaml

@@ -20,9 +20,6 @@ spec:
             accept_percentage:
               signal_name: ACCEPT_PERCENTAGE
           scheduler:
-            priority_label_key: priority
-            tokens_label_key: tokens
-            workload_label_key: workload
             workloads:
             - label_matcher:
                 match_labels:

playground/scenarios/postgres-concurrency/load-generator/test.js

@@ -0,0 +1,50 @@
+import { randomIntBetween } from "https://jslib.k6.io/k6-utils/1.4.0/index.js";
+import { check, sleep } from "k6";
+import { vu } from "k6/execution";
+import http from "k6/http";
+
+export let vuStages = [
+  { duration: "10s", target: 10 },
+  { duration: "2m", target: 50 },
+  { duration: "10s", target: 10 },
+];
+
+export let options = {
+  discardResponseBodies: true,
+  scenarios: {
+    guests: {
+      executor: "ramping-vus",
+      stages: vuStages,
+      env: { USER_TYPE: "guest" },
+    },
+    subscribers: {
+      executor: "ramping-vus",
+      stages: vuStages,
+      env: { USER_TYPE: "subscriber" },
+    },
+  },
+};
+
+export default function () {
+  let userType = __ENV.USER_TYPE;
+  let userId = vu.idInTest;
+  const url = "http://aperture-go-example.aperture-go-example.svc.cluster.local:80/postgres";
+  const headers = {
+    "Content-Type": "application/json",
+    Cookie:
+      "session=eyJ1c2VyIjoia2Vub2JpIn0.YbsY4Q.kTaKRTyOIfVlIbNB48d9YH6Q0wo",
+    "User-Type": userType,
+    "User-Id": userId,
+  };
+  const body = {}
+  let res = http.request("POST", url, JSON.stringify(body), {
+    headers: headers,
+  });
+  const ret = check(res, {
+    "http status was 200": res.status === 200,
+  });
+  if (!ret) {
+    // sleep for 10ms to 25ms
+    sleep(randomIntBetween(0.01, 0.025));
+  }
+}