Skip to content

Commit ef5cb9a

Browse files
mafredrimafredri
committed
feat(cli/exp): add app testing to scaletest workspace-traffic
1 parent 5bfbf9f commit ef5cb9a

File tree

5 files changed

+291
-22
lines changed

5 files changed

+291
-22
lines changed

‎cli/exp_scaletest.go

+39-1
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/prometheus/client_golang/prometheus"
2222
"github.com/prometheus/client_golang/prometheus/promhttp"
2323
"go.opentelemetry.io/otel/trace"
24+
"golang.org/x/exp/slices"
2425
"golang.org/x/xerrors"
2526

2627
"cdr.dev/slog"
@@ -859,6 +860,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
859860
tickInterval time.Duration
860861
bytesPerTick int64
861862
ssh bool
863+
app string
862864
template string
863865

864866
client = &codersdk.Client{}
@@ -911,6 +913,11 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
911913
}
912914
}
913915

916+
appHost, err := client.AppHost(ctx)
917+
if err != nil {
918+
return xerrors.Errorf("get app host: %w", err)
919+
}
920+
914921
workspaces, err := getScaletestWorkspaces(inv.Context(), client, template)
915922
if err != nil {
916923
return err
@@ -949,6 +956,8 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
949956
agentName string
950957
name = "workspace-traffic"
951958
id = strconv.Itoa(idx)
959+
apps []codersdk.WorkspaceApp
960+
appConfig workspacetraffic.AppConfig
952961
)
953962

954963
for _, res := range ws.LatestBuild.Resources {
@@ -957,13 +966,34 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
957966
}
958967
agentID = res.Agents[0].ID
959968
agentName = res.Agents[0].Name
969+
apps = res.Agents[0].Apps
960970
}
961971

962972
if agentID == uuid.Nil {
963973
_, _ = fmt.Fprintf(inv.Stderr, "WARN: skipping workspace %s: no agent\n", ws.Name)
964974
continue
965975
}
966976

977+
if app != "" {
978+
i := slices.IndexFunc(apps, func(a codersdk.WorkspaceApp) bool { return a.Slug == app })
979+
if i == -1 {
980+
return xerrors.Errorf("app %q not found in workspace %q", app, ws.Name)
981+
}
982+
983+
appConfig = workspacetraffic.AppConfig{
984+
Name: apps[i].Slug,
985+
}
986+
if apps[i].Subdomain {
987+
if appHost.Host == "" {
988+
return xerrors.Errorf("app %q is a subdomain app but no app host is configured", app)
989+
}
990+
991+
appConfig.URL = fmt.Sprintf("%s://%s", client.URL.Scheme, strings.Replace(appHost.Host, "*", apps[i].SubdomainName, 1))
992+
} else {
993+
appConfig.URL = fmt.Sprintf("%s/@%s/%s.%s/apps/%s", client.URL.String(), ws.OwnerName, ws.Name, agentName, apps[i].Slug)
994+
}
995+
}
996+
967997
// Setup our workspace agent connection.
968998
config := workspacetraffic.Config{
969999
AgentID: agentID,
@@ -974,6 +1004,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
9741004
WriteMetrics: metrics.WriteMetrics(ws.OwnerName, ws.Name, agentName),
9751005
SSH: ssh,
9761006
Echo: ssh,
1007+
App: appConfig,
9771008
}
9781009

9791010
if err := config.Validate(); err != nil {
@@ -1046,9 +1077,16 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
10461077
Flag: "ssh",
10471078
Env: "CODER_SCALETEST_WORKSPACE_TRAFFIC_SSH",
10481079
Default: "",
1049-
Description: "Send traffic over SSH.",
1080+
Description: "Send traffic over SSH, cannot be used with --app.",
10501081
Value: clibase.BoolOf(&ssh),
10511082
},
1083+
{
1084+
Flag: "app",
1085+
Env: "CODER_SCALETEST_WORKSPACE_TRAFFIC_APP",
1086+
Default: "",
1087+
Description: "Send WebSocket traffic to a workspace app (proxied via coderd), cannot be used with --ssh.",
1088+
Value: clibase.StringOf(&app),
1089+
},
10521090
}
10531091

10541092
tracingFlags.attach(&cmd.Options)

‎scaletest/workspacetraffic/config.go

+11
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ type Config struct {
3131
// to true will double the amount of data read from the agent for
3232
// PTYs (e.g. reconnecting pty or SSH connections that request PTY).
3333
Echo bool `json:"echo"`
34+
35+
App AppConfig `json:"app"`
3436
}
3537

3638
func (c Config) Validate() error {
@@ -50,5 +52,14 @@ func (c Config) Validate() error {
5052
return xerrors.Errorf("validate tick_interval: must be greater than zero")
5153
}
5254

55+
if c.SSH && c.App.Name != "" {
56+
return xerrors.Errorf("validate ssh: must be false when app is used")
57+
}
58+
5359
return nil
5460
}
61+
62+
type AppConfig struct {
63+
Name string `json:"name"`
64+
URL string `json:"url"`
65+
}

‎scaletest/workspacetraffic/conn.go

+118-3
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,13 @@ import (
55
"encoding/json"
66
"errors"
77
"io"
8+
"net"
9+
"net/http"
810
"sync"
911
"time"
1012

13+
"nhooyr.io/websocket"
14+
1115
"github.com/coder/coder/v2/codersdk"
1216

1317
"github.com/google/uuid"
@@ -27,9 +31,9 @@ const (
2731
//
2832
// failed to write frame: WebSocket closed: received close frame: status = StatusMessageTooBig and reason = "read limited at 32769 bytes"
2933
//
30-
// Since we can't control fragmentation/buffer sizes, we keep it simple and
31-
// match the conservative payload size used by agent/reconnectingpty (1024).
32-
rptyJSONMaxDataSize = 1024
34+
// Since we can't control fragmentation/buffer sizes, we use a conservative
35+
// value. Derived from 1024 * 9 * 3 = <28KB.
36+
rptyJSONMaxDataSize = 1024 * 9
3337
)
3438

3539
func connectRPTY(ctx context.Context, client *codersdk.Client, agentID, reconnect uuid.UUID, cmd string) (*countReadWriteCloser, error) {
@@ -260,3 +264,114 @@ func (w *wrappedSSHConn) Read(p []byte) (n int, err error) {
260264
func (w *wrappedSSHConn) Write(p []byte) (n int, err error) {
261265
return w.stdin.Write(p)
262266
}
267+
268+
func appClientConn(ctx context.Context, client *codersdk.Client, url string) (*countReadWriteCloser, error) {
269+
headers := http.Header{}
270+
tokenHeader := codersdk.SessionTokenHeader
271+
if client.SessionTokenHeader != "" {
272+
tokenHeader = client.SessionTokenHeader
273+
}
274+
headers.Set(tokenHeader, client.SessionToken())
275+
276+
//nolint:bodyclose // The websocket conn manages the body.
277+
conn, _, err := websocket.Dial(ctx, url, &websocket.DialOptions{
278+
HTTPClient: client.HTTPClient,
279+
HTTPHeader: headers,
280+
})
281+
if err != nil {
282+
return nil, xerrors.Errorf("websocket dial: %w", err)
283+
}
284+
285+
netConn := websocketNetConn(conn, websocket.MessageBinary)
286+
287+
// Wrap the conn in a countReadWriteCloser so we can monitor bytes sent/rcvd.
288+
crw := &countReadWriteCloser{rwc: netConn}
289+
return crw, nil
290+
}
291+
292+
// wsNetConn wraps net.Conn created by websocket.NetConn(). Cancel func
293+
// is called if a read or write error is encountered.
294+
type wsNetConn struct {
295+
net.Conn
296+
297+
writeMu sync.Mutex
298+
readMu sync.Mutex
299+
300+
cancel context.CancelFunc
301+
closeMu sync.Mutex
302+
closed bool
303+
}
304+
305+
func (c *wsNetConn) Read(b []byte) (n int, err error) {
306+
c.readMu.Lock()
307+
defer c.readMu.Unlock()
308+
if c.isClosed() {
309+
return 0, io.EOF
310+
}
311+
n, err = c.Conn.Read(b)
312+
if err != nil {
313+
if c.isClosed() {
314+
return n, io.EOF
315+
}
316+
return n, err
317+
}
318+
return n, nil
319+
}
320+
321+
func (c *wsNetConn) Write(b []byte) (n int, err error) {
322+
c.writeMu.Lock()
323+
defer c.writeMu.Unlock()
324+
if c.isClosed() {
325+
return 0, io.EOF
326+
}
327+
328+
for len(b) > 0 {
329+
bb := b
330+
if len(bb) > rptyJSONMaxDataSize {
331+
bb = b[:rptyJSONMaxDataSize]
332+
}
333+
b = b[len(bb):]
334+
nn, err := c.Conn.Write(bb)
335+
n += nn
336+
if err != nil {
337+
if c.isClosed() {
338+
return n, io.EOF
339+
}
340+
return n, err
341+
}
342+
}
343+
return n, nil
344+
}
345+
346+
func (c *wsNetConn) isClosed() bool {
347+
c.closeMu.Lock()
348+
defer c.closeMu.Unlock()
349+
return c.closed
350+
}
351+
352+
func (c *wsNetConn) Close() error {
353+
c.closeMu.Lock()
354+
closed := c.closed
355+
c.closed = true
356+
c.closeMu.Unlock()
357+
358+
if closed {
359+
return nil
360+
}
361+
362+
c.cancel()
363+
364+
c.readMu.Lock()
365+
defer c.readMu.Unlock()
366+
c.writeMu.Lock()
367+
defer c.writeMu.Unlock()
368+
369+
_ = c.Conn.Close()
370+
return nil
371+
}
372+
373+
func websocketNetConn(conn *websocket.Conn, msgType websocket.MessageType) net.Conn {
374+
ctx, cancel := context.WithCancel(context.Background())
375+
nc := websocket.NetConn(ctx, conn, msgType)
376+
return &wsNetConn{cancel: cancel, Conn: nc}
377+
}

‎scaletest/workspacetraffic/run.go

+22-13
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,16 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
9191
command := fmt.Sprintf("dd if=/dev/stdin of=%s bs=%d status=none", output, bytesPerTick)
9292

9393
var conn *countReadWriteCloser
94-
if r.cfg.SSH {
94+
switch {
95+
case r.cfg.App.Name != "":
96+
logger.Info(ctx, "sending traffic to workspace app", slog.F("app", r.cfg.App.Name))
97+
conn, err = appClientConn(ctx, r.client, r.cfg.App.URL)
98+
if err != nil {
99+
logger.Error(ctx, "connect to workspace app", slog.Error(err))
100+
return xerrors.Errorf("connect to workspace app: %w", err)
101+
}
102+
103+
case r.cfg.SSH:
95104
logger.Info(ctx, "connecting to workspace agent", slog.F("method", "ssh"))
96105
// If echo is enabled, disable PTY to avoid double echo and
97106
// reduce CPU usage.
@@ -101,7 +110,8 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
101110
logger.Error(ctx, "connect to workspace agent via ssh", slog.Error(err))
102111
return xerrors.Errorf("connect to workspace via ssh: %w", err)
103112
}
104-
} else {
113+
114+
default:
105115
logger.Info(ctx, "connecting to workspace agent", slog.F("method", "reconnectingpty"))
106116
conn, err = connectRPTY(ctx, r.client, agentID, reconnect, command)
107117
if err != nil {
@@ -114,8 +124,8 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
114124
closeConn := func() error {
115125
closeOnce.Do(func() {
116126
closeErr = conn.Close()
117-
if err != nil {
118-
logger.Error(ctx, "close agent connection", slog.Error(err))
127+
if closeErr != nil {
128+
logger.Error(ctx, "close agent connection", slog.Error(closeErr))
119129
}
120130
})
121131
return closeErr
@@ -142,7 +152,6 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
142152

143153
// Read until connection is closed.
144154
go func() {
145-
rch := rch // Shadowed for reassignment.
146155
logger.Debug(ctx, "reading from agent")
147156
rch <- drain(conn)
148157
logger.Debug(ctx, "done reading from agent")
@@ -151,7 +160,6 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
151160

152161
// Write random data to the conn every tick.
153162
go func() {
154-
wch := wch // Shadowed for reassignment.
155163
logger.Debug(ctx, "writing to agent")
156164
wch <- writeRandomData(conn, bytesPerTick, tick.C)
157165
logger.Debug(ctx, "done writing to agent")
@@ -160,16 +168,17 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
160168

161169
var waitCloseTimeoutCh <-chan struct{}
162170
deadlineCtxCh := deadlineCtx.Done()
171+
wchRef, rchRef := wch, rch
163172
for {
164-
if wch == nil && rch == nil {
173+
if wchRef == nil && rchRef == nil {
165174
return nil
166175
}
167176

168177
select {
169178
case <-waitCloseTimeoutCh:
170179
logger.Warn(ctx, "timed out waiting for read/write to complete",
171-
slog.F("write_done", wch == nil),
172-
slog.F("read_done", rch == nil),
180+
slog.F("write_done", wchRef == nil),
181+
slog.F("read_done", rchRef == nil),
173182
)
174183
return xerrors.Errorf("timed out waiting for read/write to complete: %w", ctx.Err())
175184
case <-deadlineCtxCh:
@@ -181,16 +190,16 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) (err error)
181190
waitCtx, cancel := context.WithTimeout(context.Background(), waitCloseTimeout)
182191
defer cancel() //nolint:revive // Only called once.
183192
waitCloseTimeoutCh = waitCtx.Done()
184-
case err = <-wch:
193+
case err = <-wchRef:
185194
if err != nil {
186195
return xerrors.Errorf("write to agent: %w", err)
187196
}
188-
wch = nil
189-
case err = <-rch:
197+
wchRef = nil
198+
case err = <-rchRef:
190199
if err != nil {
191200
return xerrors.Errorf("read from agent: %w", err)
192201
}
193-
rch = nil
202+
rchRef = nil
194203
}
195204
}
196205
}

0 commit comments

Comments
 (0)