Skip to content

Commit 88394d6

Browse files
johnstcnjohnstcn
committed
feat(cli): add trafficgen command for load testing
1 parent bb0a38b commit 88394d6

File tree

5 files changed

+311
-4
lines changed

5 files changed

+311
-4
lines changed

‎cli/root.go

+1
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ func (r *RootCmd) Core() []*clibase.Cmd {
105105
// Hidden
106106
r.workspaceAgent(),
107107
r.scaletest(),
108+
r.trafficGen(),
108109
r.gitssh(),
109110
r.vscodeSSH(),
110111
}

‎cli/scaletest_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import (
1919
)
2020

2121
func TestScaleTest(t *testing.T) {
22-
t.Skipf("This test is flakey. See https://github.com/coder/coder/issues/4942")
22+
// t.Skipf("This test is flakey. See https://github.com/coder/coder/issues/4942")
2323
t.Parallel()
2424

2525
// This test does a create-workspaces scale test with --no-cleanup, checks

‎cli/trafficgen.go

+216
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
package cli
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"github.com/coder/coder/cli/clibase"
8+
"github.com/coder/coder/codersdk"
9+
"github.com/coder/coder/cryptorand"
10+
"github.com/google/uuid"
11+
"golang.org/x/xerrors"
12+
"io"
13+
"sync/atomic"
14+
"time"
15+
)
16+
17+
func (r *RootCmd) trafficGen() *clibase.Cmd {
18+
var (
19+
duration time.Duration
20+
bps int64
21+
client = new(codersdk.Client)
22+
)
23+
24+
cmd := &clibase.Cmd{
25+
Use: "trafficgen",
26+
Hidden: true,
27+
Short: "Generate traffic to a Coder workspace",
28+
Middleware: clibase.Chain(
29+
clibase.RequireRangeArgs(1, 2),
30+
r.InitClient(client),
31+
),
32+
Handler: func(inv *clibase.Invocation) error {
33+
var agentName string
34+
ws, err := namedWorkspace(inv.Context(), client, inv.Args[0])
35+
if err != nil {
36+
return err
37+
}
38+
39+
var agentID uuid.UUID
40+
for _, res := range ws.LatestBuild.Resources {
41+
if len(res.Agents) == 0 {
42+
continue
43+
}
44+
if agentName != "" && agentName != res.Agents[0].Name {
45+
continue
46+
}
47+
agentID = res.Agents[0].ID
48+
}
49+
50+
if agentID == uuid.Nil {
51+
return xerrors.Errorf("no agent found for workspace %s", ws.Name)
52+
}
53+
54+
reconnect := uuid.New()
55+
conn, err := client.WorkspaceAgentReconnectingPTY(inv.Context(), codersdk.WorkspaceAgentReconnectingPTYOpts{
56+
AgentID: agentID,
57+
Reconnect: reconnect,
58+
Height: 65535,
59+
Width: 65535,
60+
Command: "/bin/sh",
61+
})
62+
if err != nil {
63+
return xerrors.Errorf("connect to workspace: %w", err)
64+
}
65+
66+
defer func() {
67+
_ = conn.Close()
68+
}()
69+
start := time.Now()
70+
ctx, cancel := context.WithDeadline(inv.Context(), start.Add(duration))
71+
defer cancel()
72+
crw := countReadWriter{ReadWriter: conn}
73+
// First, write a comment to the pty so we don't execute anything.
74+
data, err := json.Marshal(codersdk.ReconnectingPTYRequest{
75+
Data: "#",
76+
})
77+
if err != nil {
78+
return xerrors.Errorf("write comment to pty: %w", err)
79+
}
80+
_, err = crw.Write(data)
81+
// Now we begin writing random data to the pty.
82+
writeSize := int(bps / 10)
83+
rch := make(chan error)
84+
wch := make(chan error)
85+
go func() {
86+
rch <- readForever(ctx, &crw)
87+
close(rch)
88+
}()
89+
go func() {
90+
wch <- writeRandomData(ctx, &crw, writeSize, 100*time.Millisecond)
91+
close(wch)
92+
}()
93+
94+
if rErr := <-rch; rErr != nil {
95+
return xerrors.Errorf("read from pty: %w", rErr)
96+
}
97+
if wErr := <-wch; wErr != nil {
98+
return xerrors.Errorf("write to pty: %w", wErr)
99+
}
100+
101+
_, _ = fmt.Fprintf(inv.Stdout, "Test results:\n")
102+
_, _ = fmt.Fprintf(inv.Stdout, "Took: %.2fs\n", time.Since(start).Seconds())
103+
_, _ = fmt.Fprintf(inv.Stdout, "Sent: %d bytes\n", crw.BytesWritten())
104+
_, _ = fmt.Fprintf(inv.Stdout, "Rcvd: %d bytes\n", crw.BytesRead())
105+
return nil
106+
},
107+
}
108+
109+
cmd.Options = []clibase.Option{
110+
{
111+
Flag: "duration",
112+
Env: "CODER_TRAFFICGEN_DURATION",
113+
Default: "10s",
114+
Description: "How long to generate traffic for.",
115+
Value: clibase.DurationOf(&duration),
116+
},
117+
{
118+
Flag: "bps",
119+
Env: "CODER_TRAFFICGEN_BPS",
120+
Default: "1024",
121+
Description: "How much traffic to generate in bytes per second.",
122+
Value: clibase.Int64Of(&bps),
123+
},
124+
}
125+
126+
return cmd
127+
}
128+
129+
func readForever(ctx context.Context, src io.Reader) error {
130+
buf := make([]byte, 1024)
131+
for {
132+
select {
133+
case <-ctx.Done():
134+
return nil
135+
default:
136+
_, err := src.Read(buf)
137+
if err != nil && err != io.EOF {
138+
return err
139+
}
140+
}
141+
}
142+
}
143+
144+
func writeRandomData(ctx context.Context, dst io.Writer, size int, period time.Duration) error {
145+
tick := time.NewTicker(period)
146+
defer tick.Stop()
147+
for {
148+
select {
149+
case <-ctx.Done():
150+
return nil
151+
case <-tick.C:
152+
randStr, err := cryptorand.String(size)
153+
if err != nil {
154+
return err
155+
}
156+
data, err := json.Marshal(codersdk.ReconnectingPTYRequest{
157+
Data: randStr,
158+
})
159+
if err != nil {
160+
return err
161+
}
162+
err = copyContext(ctx, dst, data)
163+
if err != nil {
164+
return err
165+
}
166+
}
167+
}
168+
}
169+
170+
func copyContext(ctx context.Context, dst io.Writer, src []byte) error {
171+
for idx := range src {
172+
select {
173+
case <-ctx.Done():
174+
return nil
175+
default:
176+
_, err := dst.Write(src[idx : idx+1])
177+
if err != nil {
178+
if err == io.EOF {
179+
return nil
180+
}
181+
return err
182+
}
183+
}
184+
}
185+
return nil
186+
}
187+
188+
type countReadWriter struct {
189+
io.ReadWriter
190+
bytesRead atomic.Int64
191+
bytesWritten atomic.Int64
192+
}
193+
194+
func (w *countReadWriter) Read(p []byte) (int, error) {
195+
n, err := w.ReadWriter.Read(p)
196+
if err == nil {
197+
w.bytesRead.Add(int64(n))
198+
}
199+
return n, err
200+
}
201+
202+
func (w *countReadWriter) Write(p []byte) (int, error) {
203+
n, err := w.ReadWriter.Write(p)
204+
if err == nil {
205+
w.bytesWritten.Add(int64(n))
206+
}
207+
return n, err
208+
}
209+
210+
func (w *countReadWriter) BytesRead() int64 {
211+
return w.bytesRead.Load()
212+
}
213+
214+
func (w *countReadWriter) BytesWritten() int64 {
215+
return w.bytesWritten.Load()
216+
}

‎cli/trafficgen_test.go

+90
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
package cli_test
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"github.com/coder/coder/agent"
7+
"github.com/coder/coder/cli/clitest"
8+
"github.com/coder/coder/coderd/coderdtest"
9+
"github.com/coder/coder/codersdk/agentsdk"
10+
"github.com/coder/coder/provisioner/echo"
11+
"github.com/coder/coder/provisionersdk/proto"
12+
"github.com/coder/coder/testutil"
13+
"github.com/google/uuid"
14+
"github.com/stretchr/testify/require"
15+
"strings"
16+
"testing"
17+
)
18+
19+
// This test pretends to stand up a workspace and run a no-op traffic generation test.
20+
// It's not a real test, but it's useful for debugging.
21+
// We do not perform any cleanup.
22+
func TestTrafficGen(t *testing.T) {
23+
t.Parallel()
24+
25+
ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitMedium)
26+
defer cancelFunc()
27+
28+
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
29+
user := coderdtest.CreateFirstUser(t, client)
30+
31+
authToken := uuid.NewString()
32+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
33+
Parse: echo.ParseComplete,
34+
ProvisionPlan: echo.ProvisionComplete,
35+
ProvisionApply: []*proto.Provision_Response{{
36+
Type: &proto.Provision_Response_Complete{
37+
Complete: &proto.Provision_Complete{
38+
Resources: []*proto.Resource{{
39+
Name: "example",
40+
Type: "aws_instance",
41+
Agents: []*proto.Agent{{
42+
Id: uuid.NewString(),
43+
Name: "agent",
44+
Auth: &proto.Agent_Token{
45+
Token: authToken,
46+
},
47+
Apps: []*proto.App{},
48+
}},
49+
}},
50+
},
51+
},
52+
}},
53+
})
54+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
55+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
56+
57+
ws := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
58+
coderdtest.AwaitWorkspaceBuildJob(t, client, ws.LatestBuild.ID)
59+
60+
agentClient := agentsdk.New(client.URL)
61+
agentClient.SetSessionToken(authToken)
62+
agentCloser := agent.New(agent.Options{
63+
Client: agentClient,
64+
})
65+
t.Cleanup(func() {
66+
_ = agentCloser.Close()
67+
})
68+
69+
coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
70+
71+
inv, root := clitest.New(t, "trafficgen", ws.Name,
72+
"--duration", "1s",
73+
"--bps", "100",
74+
)
75+
clitest.SetupConfig(t, client, root)
76+
var stdout, stderr bytes.Buffer
77+
inv.Stdout = &stdout
78+
inv.Stderr = &stderr
79+
err := inv.WithContext(ctx).Run()
80+
require.NoError(t, err)
81+
stdoutStr := stdout.String()
82+
stderrStr := stderr.String()
83+
require.Empty(t, stderrStr)
84+
lines := strings.Split(strings.TrimSpace(stdoutStr), "\n")
85+
require.Len(t, lines, 4)
86+
require.Equal(t, "Test results:", lines[0])
87+
require.Regexp(t, `Took:\s+\d+\.\d+s`, lines[1])
88+
require.Regexp(t, `Sent:\s+\d+ bytes`, lines[2])
89+
require.Regexp(t, `Rcvd:\s+\d+ bytes`, lines[3])
90+
}

‎codersdk/workspaceagentconn.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -165,9 +165,9 @@ type WorkspaceAgentReconnectingPTYInit struct {
165165
// to pipe data to a PTY.
166166
// @typescript-ignore ReconnectingPTYRequest
167167
type ReconnectingPTYRequest struct {
168-
Data string `json:"data"`
169-
Height uint16 `json:"height"`
170-
Width uint16 `json:"width"`
168+
Data string `json:"data,omitempty"`
169+
Height uint16 `json:"height,omitempty"`
170+
Width uint16 `json:"width,omitempty"`
171171
}
172172

173173
// ReconnectingPTY spawns a new reconnecting terminal session.

0 commit comments

Comments
 (0)