feat: add workspace-proxy-url flag to scaletest workspace-traffic

f0ssel · f0ssel · commit 55e1ec90cbab · 2024-12-18T16:29:52.000Z
diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"math/rand"
 	"net/http"
+	"net/url"
 	"os"
 	"os/signal"
 	"strconv"
@@ -860,13 +861,14 @@ func (r *RootCmd) scaletestCreateWorkspaces() *serpent.Command {
 
 func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 	var (
-		tickInterval     time.Duration
-		bytesPerTick     int64
-		ssh              bool
-		useHostLogin     bool
-		app              string
-		template         string
-		targetWorkspaces string
+		tickInterval      time.Duration
+		bytesPerTick      int64
+		ssh               bool
+		useHostLogin      bool
+		app               string
+		template          string
+		targetWorkspaces  string
+		workspaceProxyURL string
 
 		client          = &codersdk.Client{}
 		tracingFlags    = &scaletestTracingFlags{}
@@ -1002,6 +1004,23 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 					return xerrors.Errorf("configure workspace app: %w", err)
 				}
 
+				var webClient *codersdk.Client
+				if workspaceProxyURL != "" {
+					u, err := url.Parse(workspaceProxyURL)
+					if err != nil {
+						return xerrors.Errorf("parse workspace proxy URL: %w", err)
+					}
+
+					webClient = codersdk.New(u)
+					webClient.HTTPClient = client.HTTPClient
+					webClient.SetSessionToken(client.SessionToken())
+
+					appConfig, err = createWorkspaceAppConfig(webClient, appHost.Host, app, ws, agent)
+					if err != nil {
+						return xerrors.Errorf("configure proxy workspace app: %w", err)
+					}
+				}
+
 				// Setup our workspace agent connection.
 				config := workspacetraffic.Config{
 					AgentID:      agent.ID,
@@ -1015,6 +1034,10 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 					App:          appConfig,
 				}
 
+				if webClient != nil {
+					config.WebClient = webClient
+				}
+
 				if err := config.Validate(); err != nil {
 					return xerrors.Errorf("validate config: %w", err)
 				}
@@ -1108,6 +1131,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *serpent.Command {
 			Description: "Connect as the currently logged in user.",
 			Value:       serpent.BoolOf(&useHostLogin),
 		},
+		{
+			Flag:        "workspace-proxy-url",
+			Env:         "CODER_SCALETEST_WORKSPACE_PROXY_URL",
+			Default:     "",
+			Description: "URL for workspace proxy to send web traffic to.",
+			Value:       serpent.StringOf(&workspaceProxyURL),
+		},
 	}
 
 	tracingFlags.attach(&cmd.Options)
diff --git a/scaletest/workspacetraffic/config.go b/scaletest/workspacetraffic/config.go
@@ -5,6 +5,8 @@ import (
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/codersdk"
 )
 
 type Config struct {
@@ -33,6 +35,8 @@ type Config struct {
 	Echo bool `json:"echo"`
 
 	App AppConfig `json:"app"`
+
+	WebClient *codersdk.Client
 }
 
 func (c Config) Validate() error {
diff --git a/scaletest/workspacetraffic/run.go b/scaletest/workspacetraffic/run.go
@@ -23,8 +23,9 @@ import (
 )
 
 type Runner struct {
-	client *codersdk.Client
-	cfg    Config
+	client    *codersdk.Client
+	webClient *codersdk.Client
+	cfg       Config
 }
 
 var (
@@ -34,9 +35,15 @@ var (
 
 // func NewRunner(client *codersdk.Client, cfg Config, metrics *Metrics) *Runner {
 func NewRunner(client *codersdk.Client, cfg Config) *Runner {
+	webClient := client
+	if cfg.WebClient != nil {
+		webClient = cfg.WebClient
+	}
+
 	return &Runner{
-		client: client,
-		cfg:    cfg,
+		client:    client,
+		webClient: webClient,
+		cfg:       cfg,
 	}
 }
 
@@ -94,7 +101,7 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
 	switch {
 	case r.cfg.App.Name != "":
 		logger.Info(ctx, "sending traffic to workspace app", slog.F("app", r.cfg.App.Name))
-		conn, err = appClientConn(ctx, r.client, r.cfg.App.URL)
+		conn, err = appClientConn(ctx, r.webClient, r.cfg.App.URL)
 		if err != nil {
 			logger.Error(ctx, "connect to workspace app", slog.Error(err))
 			return xerrors.Errorf("connect to workspace app: %w", err)
@@ -113,7 +120,7 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
 
 	default:
 		logger.Info(ctx, "connecting to workspace agent", slog.F("method", "reconnectingpty"))
-		conn, err = connectRPTY(ctx, r.client, agentID, reconnect, command)
+		conn, err = connectRPTY(ctx, r.webClient, agentID, reconnect, command)
 		if err != nil {
 			logger.Error(ctx, "connect to workspace agent via reconnectingpty", slog.Error(err))
 			return xerrors.Errorf("connect to workspace via reconnectingpty: %w", err)

Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,8 @@ import (`
`5`	`5`
`6`	`6`	`"github.com/google/uuid"`
`7`	`7`	`"golang.org/x/xerrors"`
	`8`	`+`
	`9`	`+ "github.com/coder/coder/v2/codersdk"`
`8`	`10`	`)`
`9`	`11`
`10`	`12`	`type Config struct {`
`@@ -33,6 +35,8 @@ type Config struct {`
`33`	`35`	Echo bool `json:"echo"`
`34`	`36`
`35`	`37`	App AppConfig `json:"app"`
	`38`	`+`
	`39`	`+ WebClient *codersdk.Client`
`36`	`40`	`}`
`37`	`41`
`38`	`42`	`func (c Config) Validate() error {`