@@ -14,9 +14,14 @@ import (
14
14
"time"
15
15
16
16
"github.com/google/uuid"
17
+ "github.com/prometheus/client_golang/prometheus"
18
+ "github.com/prometheus/client_golang/prometheus/promhttp"
17
19
"go.opentelemetry.io/otel/trace"
18
20
"golang.org/x/xerrors"
19
21
22
+ "cdr.dev/slog"
23
+ "cdr.dev/slog/sloggers/sloghuman"
24
+
20
25
"github.com/coder/coder/cli/clibase"
21
26
"github.com/coder/coder/cli/cliui"
22
27
"github.com/coder/coder/coderd/httpapi"
@@ -896,8 +901,11 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
896
901
897
902
func (r * RootCmd ) scaletestWorkspaceTraffic () * clibase.Cmd {
898
903
var (
899
- tickInterval time.Duration
900
- bytesPerTick int64
904
+ tickInterval time.Duration
905
+ bytesPerTick int64
906
+ scaletestPrometheusAddress string
907
+ scaletestPrometheusWait time.Duration
908
+
901
909
client = & codersdk.Client {}
902
910
tracingFlags = & scaletestTracingFlags {}
903
911
strategy = & scaletestStrategyFlags {}
@@ -913,6 +921,12 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
913
921
),
914
922
Handler : func (inv * clibase.Invocation ) error {
915
923
ctx := inv .Context ()
924
+ reg := prometheus .NewRegistry ()
925
+ metrics := workspacetraffic .NewMetrics (reg , "username" , "workspace_name" , "agent_name" )
926
+
927
+ logger := slog .Make (sloghuman .Sink (io .Discard ))
928
+ prometheusSrvClose := ServeHandler (ctx , logger , promhttp .HandlerFor (reg , promhttp.HandlerOpts {}), scaletestPrometheusAddress , "prometheus" )
929
+ defer prometheusSrvClose ()
916
930
917
931
// Bypass rate limiting
918
932
client .HTTPClient = & http.Client {
@@ -943,6 +957,9 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
943
957
_ , _ = fmt .Fprintln (inv .Stderr , "\n Uploading traces..." )
944
958
if err := closeTracing (ctx ); err != nil {
945
959
_ , _ = fmt .Fprintf (inv .Stderr , "\n Error uploading traces: %+v\n " , err )
960
+ // Wait for prometheus metrics to be scraped
961
+ _ , _ = fmt .Fprintf (inv .Stderr , "Waiting %s for prometheus metrics to be scraped\n " , scaletestPrometheusWait )
962
+ <- time .After (scaletestPrometheusWait )
946
963
}
947
964
}()
948
965
tracer := tracerProvider .Tracer (scaletestTracerName )
@@ -955,16 +972,18 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
955
972
th := harness .NewTestHarness (strategy .toStrategy (), cleanupStrategy .toStrategy ())
956
973
for idx , ws := range workspaces {
957
974
var (
958
- agentID uuid.UUID
959
- name = "workspace-traffic"
960
- id = strconv .Itoa (idx )
975
+ agentID uuid.UUID
976
+ agentName string
977
+ name = "workspace-traffic"
978
+ id = strconv .Itoa (idx )
961
979
)
962
980
963
981
for _ , res := range ws .LatestBuild .Resources {
964
982
if len (res .Agents ) == 0 {
965
983
continue
966
984
}
967
985
agentID = res .Agents [0 ].ID
986
+ agentName = res .Agents [0 ].Name
968
987
}
969
988
970
989
if agentID == uuid .Nil {
@@ -974,16 +993,20 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
974
993
975
994
// Setup our workspace agent connection.
976
995
config := workspacetraffic.Config {
977
- AgentID : agentID ,
978
- BytesPerTick : bytesPerTick ,
979
- Duration : strategy .timeout ,
980
- TickInterval : tickInterval ,
996
+ AgentID : agentID ,
997
+ AgentName : agentName ,
998
+ BytesPerTick : bytesPerTick ,
999
+ Duration : strategy .timeout ,
1000
+ TickInterval : tickInterval ,
1001
+ WorkspaceName : ws .Name ,
1002
+ WorkspaceOwner : ws .OwnerName ,
1003
+ Registry : reg ,
981
1004
}
982
1005
983
1006
if err := config .Validate (); err != nil {
984
1007
return xerrors .Errorf ("validate config: %w" , err )
985
1008
}
986
- var runner harness.Runnable = workspacetraffic .NewRunner (client , config )
1009
+ var runner harness.Runnable = workspacetraffic .NewRunner (client , config , metrics )
987
1010
if tracingEnabled {
988
1011
runner = & runnableTraceWrapper {
989
1012
tracer : tracer ,
@@ -1034,6 +1057,20 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
1034
1057
Description : "How often to send traffic." ,
1035
1058
Value : clibase .DurationOf (& tickInterval ),
1036
1059
},
1060
+ {
1061
+ Flag : "scaletest-prometheus-address" ,
1062
+ Env : "CODER_SCALETEST_PROMETHEUS_ADDRESS" ,
1063
+ Default : "0.0.0.0:21112" ,
1064
+ Description : "Address on which to expose scaletest Prometheus metrics." ,
1065
+ Value : clibase .StringOf (& scaletestPrometheusAddress ),
1066
+ },
1067
+ {
1068
+ Flag : "scaletest-prometheus-wait" ,
1069
+ Env : "CODER_SCALETEST_PROMETHEUS_WAIT" ,
1070
+ Default : "5s" ,
1071
+ Description : "How long to wait before exiting in order to allow Prometheus metrics to be scraped." ,
1072
+ Value : clibase .DurationOf (& scaletestPrometheusWait ),
1073
+ },
1037
1074
}
1038
1075
1039
1076
tracingFlags .attach (& cmd .Options )
0 commit comments