@@ -4,25 +4,33 @@ package cli
44
55import (
66"fmt"
7+ "net/http"
8+ "time"
79
810"github.com/prometheus/client_golang/prometheus"
11+ "github.com/prometheus/client_golang/prometheus/promhttp"
912"golang.org/x/xerrors"
1013
1114"cdr.dev/slog"
1215"cdr.dev/slog/sloggers/sloghuman"
16+ "github.com/coder/serpent"
1317
18+ "github.com/coder/coder/v2/codersdk"
1419"github.com/coder/coder/v2/scaletest/dynamicparameters"
1520"github.com/coder/coder/v2/scaletest/harness"
16- "github.com/coder/serpent"
1721)
1822
1923const (
2024dynamicParametersTestName = "dynamic-parameters"
2125)
2226
2327func (r * RootCmd )scaletestDynamicParameters ()* serpent.Command {
24- var templateName string
25- var numEvals int64
28+ var (
29+ templateName string
30+ numEvals int64
31+ tracingFlags = & scaletestTracingFlags {}
32+ prometheusFlags = & scaletestPrometheusFlags {}
33+ )
2634orgContext := NewOrganizationContext ()
2735output := & scaletestOutputFlags {}
2836
@@ -51,15 +59,50 @@ func (r *RootCmd) scaletestDynamicParameters() *serpent.Command {
5159return err
5260}
5361
62+ _ ,err = requireAdmin (ctx ,client )
63+ if err != nil {
64+ return err
65+ }
66+
67+ client .HTTPClient = & http.Client {
68+ Transport :& codersdk.HeaderTransport {
69+ Transport :http .DefaultTransport ,
70+ Header :map [string ][]string {
71+ codersdk .BypassRatelimitHeader : {"true" },
72+ },
73+ },
74+ }
75+
76+ reg := prometheus .NewRegistry ()
77+ metrics := dynamicparameters .NewMetrics (reg ,"concurrent_evaluations" )
78+
5479logger := slog .Make (sloghuman .Sink (inv .Stdout )).Leveled (slog .LevelDebug )
80+ prometheusSrvClose := ServeHandler (ctx ,logger ,promhttp .HandlerFor (reg , promhttp.HandlerOpts {}),prometheusFlags .Address ,"prometheus" )
81+ defer prometheusSrvClose ()
82+
83+ tracerProvider ,closeTracing ,tracingEnabled ,err := tracingFlags .provider (ctx )
84+ if err != nil {
85+ return xerrors .Errorf ("create tracer provider: %w" ,err )
86+ }
87+ defer func () {
88+ // Allow time for traces to flush even if command context is
89+ // canceled. This is a no-op if tracing is not enabled.
90+ _ ,_ = fmt .Fprintln (inv .Stderr ,"\n Uploading traces..." )
91+ if err := closeTracing (ctx );err != nil {
92+ _ ,_ = fmt .Fprintf (inv .Stderr ,"\n Error uploading traces: %+v\n " ,err )
93+ }
94+ // Wait for prometheus metrics to be scraped
95+ _ ,_ = fmt .Fprintf (inv .Stderr ,"Waiting %s for prometheus metrics to be scraped\n " ,prometheusFlags .Wait )
96+ <- time .After (prometheusFlags .Wait )
97+ }()
98+ tracer := tracerProvider .Tracer (scaletestTracerName )
99+
55100partitions ,err := dynamicparameters .SetupPartitions (ctx ,client ,org .ID ,templateName ,numEvals ,logger )
56101if err != nil {
57102return xerrors .Errorf ("setup dynamic parameters partitions: %w" ,err )
58103}
59104
60105th := harness .NewTestHarness (harness.ConcurrentExecutionStrategy {}, harness.ConcurrentExecutionStrategy {})
61- reg := prometheus .NewRegistry ()
62- metrics := dynamicparameters .NewMetrics (reg ,"concurrent_evaluations" )
63106
64107for i ,part := range partitions {
65108for j := range part .ConcurrentEvaluations {
@@ -68,7 +111,14 @@ func (r *RootCmd) scaletestDynamicParameters() *serpent.Command {
68111Metrics :metrics ,
69112MetricLabelValues : []string {fmt .Sprintf ("%d" ,part .ConcurrentEvaluations )},
70113}
71- runner := dynamicparameters .NewRunner (client ,cfg )
114+ var runner harness.Runnable = dynamicparameters .NewRunner (client ,cfg )
115+ if tracingEnabled {
116+ runner = & runnableTraceWrapper {
117+ tracer :tracer ,
118+ spanName :fmt .Sprintf ("%s/%d/%d" ,dynamicParametersTestName ,i ,j ),
119+ runner :runner ,
120+ }
121+ }
72122th .AddRun (dynamicParametersTestName ,fmt .Sprintf ("%d/%d" ,j ,i ),runner )
73123}
74124}
@@ -106,5 +156,7 @@ func (r *RootCmd) scaletestDynamicParameters() *serpent.Command {
106156}
107157orgContext .AttachOptions (cmd )
108158output .attach (& cmd .Options )
159+ tracingFlags .attach (& cmd .Options )
160+ prometheusFlags .attach (& cmd .Options )
109161return cmd
110162}