From 095572999acc734fc331ae83bb50ef4f9d8a1e94 Mon Sep 17 00:00:00 2001
From: Oliver Tan <otan@cockroachlabs.com>
Date: Fri, 24 Sep 2021 18:54:19 +1000
Subject: [PATCH] democluster: expose flag for workload rate limit on demo

Release note (cli change): The 25 max QPS rate limit for workloads on
`cockroach` demo can now be configured with a `workload-max-qps` flag.
---
 pkg/cli/cliflags/flags.go           | 5 +++++
 pkg/cli/context.go                  | 1 +
 pkg/cli/democluster/context.go      | 4 ++++
 pkg/cli/democluster/demo_cluster.go | 4 ++--
 pkg/cli/flags.go                    | 1 +
 5 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/pkg/cli/cliflags/flags.go b/pkg/cli/cliflags/flags.go
index e9f53d45299b..49e3c133dce3 100644
--- a/pkg/cli/cliflags/flags.go
+++ b/pkg/cli/cliflags/flags.go
@@ -1169,6 +1169,11 @@ can also be specified (e.g. .25).`,
 		Description: `Run a demo workload against the pre-loaded database.`,
 	}
 
+	DemoWorkloadMaxQPS = FlagInfo{
+		Name:        "workload-max-qps",
+		Description: "The maximum QPS when a workload is running.",
+	}
+
 	DemoNodeLocality = FlagInfo{
 		Name: "demo-locality",
 		Description: `
diff --git a/pkg/cli/context.go b/pkg/cli/context.go
index 1249a361b206..8b3482ef15f2 100644
--- a/pkg/cli/context.go
+++ b/pkg/cli/context.go
@@ -584,6 +584,7 @@ func setDemoContextDefaults() {
 	demoCtx.Insecure = false
 	demoCtx.SQLPort, _ = strconv.Atoi(base.DefaultPort)
 	demoCtx.HTTPPort, _ = strconv.Atoi(base.DefaultHTTPPort)
+	demoCtx.WorkloadMaxQPS = 25
 }
 
 // stmtDiagCtx captures the command-line parameters of the 'statement-diag'
diff --git a/pkg/cli/democluster/context.go b/pkg/cli/democluster/context.go
index d44432bac007..711a3394de0d 100644
--- a/pkg/cli/democluster/context.go
+++ b/pkg/cli/democluster/context.go
@@ -53,6 +53,10 @@ type Context struct {
 	// WorkloadGenerator is the desired workload generator.
 	WorkloadGenerator workload.Generator
 
+	// WorkloadMaxQPS controls the amount of queries that can be run per
+	// second.
+	WorkloadMaxQPS int
+
 	// Localities configures the list of localities available for use
 	// by instantiated servers.
 	Localities DemoLocalityList
diff --git a/pkg/cli/democluster/demo_cluster.go b/pkg/cli/democluster/demo_cluster.go
index 352316d272d2..7f76c4c507a9 100644
--- a/pkg/cli/democluster/demo_cluster.go
+++ b/pkg/cli/democluster/demo_cluster.go
@@ -1091,8 +1091,8 @@ func (c *transientCluster) runWorkload(
 		return errors.Wrap(err, "unable to create workload")
 	}
 
-	// Use a light rate limit of 25 queries per second
-	limiter := rate.NewLimiter(rate.Limit(25), 1)
+	// Use a rate limit (default 25 queries per second).
+	limiter := rate.NewLimiter(rate.Limit(c.demoCtx.WorkloadMaxQPS), 1)
 
 	// Start a goroutine to run each of the workload functions.
 	for _, workerFn := range ops.WorkerFns {
diff --git a/pkg/cli/flags.go b/pkg/cli/flags.go
index dca18403d6e4..e85d2d7cb5b0 100644
--- a/pkg/cli/flags.go
+++ b/pkg/cli/flags.go
@@ -811,6 +811,7 @@ func init() {
 
 		intFlag(f, &demoCtx.NumNodes, cliflags.DemoNodes)
 		boolFlag(f, &demoCtx.RunWorkload, cliflags.RunDemoWorkload)
+		intFlag(f, &demoCtx.WorkloadMaxQPS, cliflags.DemoWorkloadMaxQPS)
 		varFlag(f, &demoCtx.Localities, cliflags.DemoNodeLocality)
 		boolFlag(f, &demoCtx.GeoPartitionedReplicas, cliflags.DemoGeoPartitionedReplicas)
 		varFlag(f, demoNodeSQLMemSizeValue, cliflags.DemoNodeSQLMemSize)