feat: add workspace-proxy-url flag to scaletest workspace-traffic (#15920)

This allows the command to target a workspace proxy when appropriate. 

Part of https://github.com/coder/internal/issues/149

So far I've verified the correct url being used with logs:
```
➜  coder git:(f0ssel/workspace-traffic-proxy) ✗ go run ./cmd/coder/main.go exp scaletest workspace-traffic --job-timeout=60s --workspace-proxy-url="https://paris.fly.dev.coder.com"
Running load test...

web url: https://paris.fly.dev.coder.com

...

➜  coder git:(f0ssel/workspace-traffic-proxy) ✗ go run ./cmd/coder/main.go exp scaletest workspace-traffic --job-timeout=60s --workspace-proxy-url="https://paris.fly.dev.coder.com" --app="code-server"
Running load test...

app url: https://paris.fly.dev.coder.com/@f0ssel/scaletest-1.dev/apps/code-server
```
This commit is contained in:
Garrett Delfosse
2024-12-18 11:53:01 -05:00
committed by GitHub
parent 7be96bbb09
commit 50333d312f
3 changed files with 54 additions and 13 deletions
+37 -7
View File
@@ -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)
+4
View File
@@ -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 {
+13 -6
View File
@@ -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)