Merge #70642

70642: democluster: expose flag for workload rate limit on demo r=knz a=otan 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. Co-authored-by: Oliver Tan <[email protected]>
cockroachdb · Sep 24, 2021 · 6994559 · 6994559
2 parents 117de30 + 0955729
commit 6994559
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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)