Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit2ffd742

Browse files
committed
feat(cli): add trafficgen command for load testing
1 parent465fe86 commit2ffd742

File tree

5 files changed

+311
-4
lines changed

5 files changed

+311
-4
lines changed

‎cli/root.go

Lines changed: 1 addition & 0 deletions
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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import (
1919
)
2020

2121
funcTestScaleTest(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

Lines changed: 216 additions & 0 deletions
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+
bpsint64
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+
varagentNamestring
34+
ws,err:=namedWorkspace(inv.Context(),client,inv.Args[0])
35+
iferr!=nil {
36+
returnerr
37+
}
38+
39+
varagentID uuid.UUID
40+
for_,res:=rangews.LatestBuild.Resources {
41+
iflen(res.Agents)==0 {
42+
continue
43+
}
44+
ifagentName!=""&&agentName!=res.Agents[0].Name {
45+
continue
46+
}
47+
agentID=res.Agents[0].ID
48+
}
49+
50+
ifagentID==uuid.Nil {
51+
returnxerrors.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+
iferr!=nil {
63+
returnxerrors.Errorf("connect to workspace: %w",err)
64+
}
65+
66+
deferfunc() {
67+
_=conn.Close()
68+
}()
69+
start:=time.Now()
70+
ctx,cancel:=context.WithDeadline(inv.Context(),start.Add(duration))
71+
defercancel()
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+
iferr!=nil {
78+
returnxerrors.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(chanerror)
84+
wch:=make(chanerror)
85+
gofunc() {
86+
rch<-readForever(ctx,&crw)
87+
close(rch)
88+
}()
89+
gofunc() {
90+
wch<-writeRandomData(ctx,&crw,writeSize,100*time.Millisecond)
91+
close(wch)
92+
}()
93+
94+
ifrErr:=<-rch;rErr!=nil {
95+
returnxerrors.Errorf("read from pty: %w",rErr)
96+
}
97+
ifwErr:=<-wch;wErr!=nil {
98+
returnxerrors.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+
returnnil
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+
returncmd
127+
}
128+
129+
funcreadForever(ctx context.Context,src io.Reader)error {
130+
buf:=make([]byte,1024)
131+
for {
132+
select {
133+
case<-ctx.Done():
134+
returnnil
135+
default:
136+
_,err:=src.Read(buf)
137+
iferr!=nil&&err!=io.EOF {
138+
returnerr
139+
}
140+
}
141+
}
142+
}
143+
144+
funcwriteRandomData(ctx context.Context,dst io.Writer,sizeint,period time.Duration)error {
145+
tick:=time.NewTicker(period)
146+
defertick.Stop()
147+
for {
148+
select {
149+
case<-ctx.Done():
150+
returnnil
151+
case<-tick.C:
152+
randStr,err:=cryptorand.String(size)
153+
iferr!=nil {
154+
returnerr
155+
}
156+
data,err:=json.Marshal(codersdk.ReconnectingPTYRequest{
157+
Data:randStr,
158+
})
159+
iferr!=nil {
160+
returnerr
161+
}
162+
err=copyContext(ctx,dst,data)
163+
iferr!=nil {
164+
returnerr
165+
}
166+
}
167+
}
168+
}
169+
170+
funccopyContext(ctx context.Context,dst io.Writer,src []byte)error {
171+
foridx:=rangesrc {
172+
select {
173+
case<-ctx.Done():
174+
returnnil
175+
default:
176+
_,err:=dst.Write(src[idx :idx+1])
177+
iferr!=nil {
178+
iferr==io.EOF {
179+
returnnil
180+
}
181+
returnerr
182+
}
183+
}
184+
}
185+
returnnil
186+
}
187+
188+
typecountReadWriterstruct {
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+
iferr==nil {
197+
w.bytesRead.Add(int64(n))
198+
}
199+
returnn,err
200+
}
201+
202+
func (w*countReadWriter)Write(p []byte) (int,error) {
203+
n,err:=w.ReadWriter.Write(p)
204+
iferr==nil {
205+
w.bytesWritten.Add(int64(n))
206+
}
207+
returnn,err
208+
}
209+
210+
func (w*countReadWriter)BytesRead()int64 {
211+
returnw.bytesRead.Load()
212+
}
213+
214+
func (w*countReadWriter)BytesWritten()int64 {
215+
returnw.bytesWritten.Load()
216+
}

‎cli/trafficgen_test.go

Lines changed: 90 additions & 0 deletions
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+
funcTestTrafficGen(t*testing.T) {
23+
t.Parallel()
24+
25+
ctx,cancelFunc:=context.WithTimeout(context.Background(),testutil.WaitMedium)
26+
defercancelFunc()
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+
varstdout,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

Lines changed: 3 additions & 3 deletions
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
typeReconnectingPTYRequeststruct {
168-
Datastring`json:"data"`
169-
Heightuint16`json:"height"`
170-
Widthuint16`json:"width"`
168+
Datastring`json:"data,omitempty"`
169+
Heightuint16`json:"height,omitempty"`
170+
Widthuint16`json:"width,omitempty"`
171171
}
172172

173173
// ReconnectingPTY spawns a new reconnecting terminal session.

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp