add wait for prometheus metrics

coder · johnstcn · May 26, 2023 · May 15, 2023 · May 16, 2023 · May 16, 2023
commit 239ba960706431261e1bd6f24b6f2f97a5b1d64c
diff --git a/cli/scaletest.go b/cli/scaletest.go
@@ -904,6 +904,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
 		tickInterval               time.Duration
 		bytesPerTick               int64
 		scaletestPrometheusAddress string
+		scaletestPrometheusWait    time.Duration
 
 		client          = &codersdk.Client{}
 		tracingFlags    = &scaletestTracingFlags{}
@@ -956,6 +957,9 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
 				_, _ = fmt.Fprintln(inv.Stderr, "\nUploading traces...")
 				if err := closeTracing(ctx); err != nil {
 					_, _ = fmt.Fprintf(inv.Stderr, "\nError uploading traces: %+v\n", err)
+					// Wait for prometheus metrics to be scraped
+					_, _ = fmt.Fprintf(inv.Stderr, "Waiting %s for prometheus metrics to be scraped\n", scaletestPrometheusWait)
+					<-time.After(scaletestPrometheusWait)
 				}
 			}()
 			tracer := tracerProvider.Tracer(scaletestTracerName)
@@ -1060,6 +1064,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
 			Description: "Address on which to expose scaletest Prometheus metrics.",
 			Value:       clibase.StringOf(&scaletestPrometheusAddress),
 		},
+		{
+			Flag:        "scaletest-prometheus-wait",
+			Env:         "CODER_SCALETEST_PROMETHEUS_WAIT",
+			Default:     "5s",
+			Description: "How long to wait before exiting in order to allow Prometheus metrics to be scraped.",
+			Value:       clibase.DurationOf(&scaletestPrometheusWait),
+		},
 	}
 
 	tracingFlags.attach(&cmd.Options)

diff --git a/cli/scaletest_test.go b/cli/scaletest_test.go
@@ -216,6 +216,7 @@ func TestScaleTestWorkspaceTraffic(t *testing.T) {
 		"--bytes-per-tick", "1024",
 		"--tick-interval", "100ms",
 		"--scaletest-prometheus-address", "127.0.0.1:0",
+		"--scaletest-prometheus-wait", "0s",
 	)
 	clitest.SetupConfig(t, client, root)
 	var stdout, stderr bytes.Buffer

diff --git a/cli/testdata/coder_scaletest_workspace-traffic_--help.golden b/cli/testdata/coder_scaletest_workspace-traffic_--help.golden
@@ -30,6 +30,10 @@ Generate traffic to scaletest workspaces through coderd
       --scaletest-prometheus-address string, $CODER_SCALETEST_PROMETHEUS_ADDRESS (default: 0.0.0.0:21112)
           Address on which to expose scaletest Prometheus metrics.
 
+      --scaletest-prometheus-wait duration, $CODER_SCALETEST_PROMETHEUS_WAIT (default: 5s)
+          How long to wait before exiting in order to allow Prometheus metrics
+          to be scraped.
+
       --tick-interval duration, $CODER_SCALETEST_WORKSPACE_TRAFFIC_TICK_INTERVAL (default: 100ms)
           How often to send traffic.
 

diff --git a/docs/cli/scaletest_workspace-traffic.md b/docs/cli/scaletest_workspace-traffic.md
@@ -92,6 +92,16 @@ Output format specs in the format "<format>[:<path>]". Not specifying a path wil
 
 Address on which to expose scaletest Prometheus metrics.
 
+### --scaletest-prometheus-wait
+
+|             |                                               |
+| ----------- | --------------------------------------------- |
+| Type        | <code>duration</code>                         |
+| Environment | <code>$CODER_SCALETEST_PROMETHEUS_WAIT</code> |
+| Default     | <code>5s</code>                               |
+
+How long to wait before exiting in order to allow Prometheus metrics to be scraped.
+
 ### --tick-interval
 
 |             |                                                               |