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

Commitb0a641d

Browse files
committed
close connection, add output formatting
1 parent04bc551 commitb0a641d

File tree

2 files changed

+112
-56
lines changed

2 files changed

+112
-56
lines changed

‎cli/trafficgen.go

Lines changed: 100 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,32 @@ import (
1212
"golang.org/x/xerrors"
1313

1414
"github.com/coder/coder/cli/clibase"
15+
"github.com/coder/coder/cli/cliui"
1516
"github.com/coder/coder/codersdk"
1617
"github.com/coder/coder/cryptorand"
1718
)
1819

20+
typetrafficGenOutputstruct {
21+
DurationSecondsfloat64`json:"duration_s"`
22+
SentBytesint64`json:"sent_bytes"`
23+
RcvdBytesint64`json:"rcvd_bytes"`
24+
}
25+
26+
func (otrafficGenOutput)String()string {
27+
returnfmt.Sprintf("Duration: %.2fs\n",o.DurationSeconds)+
28+
fmt.Sprintf("Sent: %dB\n",o.SentBytes)+
29+
fmt.Sprintf("Rcvd: %dB",o.RcvdBytes)
30+
}
31+
1932
func (r*RootCmd)trafficGen()*clibase.Cmd {
2033
var (
21-
duration time.Duration
22-
bpsint64
23-
client=new(codersdk.Client)
34+
duration time.Duration
35+
formatter=cliui.NewOutputFormatter(
36+
cliui.TextFormat(),
37+
cliui.JSONFormat(),
38+
)
39+
bpsint64
40+
client=new(codersdk.Client)
2441
)
2542

2643
cmd:=&clibase.Cmd{
@@ -32,7 +49,10 @@ func (r *RootCmd) trafficGen() *clibase.Cmd {
3249
r.InitClient(client),
3350
),
3451
Handler:func(inv*clibase.Invocation)error {
35-
varagentNamestring
52+
var (
53+
agentNamestring
54+
tickInterval=100*time.Millisecond
55+
)
3656
ws,err:=namedWorkspace(inv.Context(),client,inv.Args[0])
3757
iferr!=nil {
3858
returnerr
@@ -53,6 +73,7 @@ func (r *RootCmd) trafficGen() *clibase.Cmd {
5373
returnxerrors.Errorf("no agent found for workspace %s",ws.Name)
5474
}
5575

76+
// Setup our workspace agent connection.
5677
reconnect:=uuid.New()
5778
conn,err:=client.WorkspaceAgentReconnectingPTY(inv.Context(), codersdk.WorkspaceAgentReconnectingPTYOpts{
5879
AgentID:agentID,
@@ -68,46 +89,60 @@ func (r *RootCmd) trafficGen() *clibase.Cmd {
6889
deferfunc() {
6990
_=conn.Close()
7091
}()
92+
93+
// Wrap the conn in a countReadWriter so we can monitor bytes sent/rcvd.
94+
crw:=countReadWriter{ReadWriter:conn}
95+
96+
// Set a deadline for stopping the text.
7197
start:=time.Now()
72-
ctx,cancel:=context.WithDeadline(inv.Context(),start.Add(duration))
98+
deadlineCtx,cancel:=context.WithDeadline(inv.Context(),start.Add(duration))
7399
defercancel()
74-
crw:=countReadWriter{ReadWriter:conn}
75-
// First, write a comment to the pty so we don't execute anything.
76-
data,err:=json.Marshal(codersdk.ReconnectingPTYRequest{
77-
Data:"#",
78-
})
79-
iferr!=nil {
80-
returnxerrors.Errorf("serialize request: %w",err)
81-
}
82-
_,err=crw.Write(data)
83-
iferr!=nil {
84-
returnxerrors.Errorf("write comment to pty: %w",err)
85-
}
100+
101+
// Create a ticker for sending data to the PTY.
102+
tick:=time.NewTicker(tickInterval)
103+
defertick.Stop()
104+
86105
// Now we begin writing random data to the pty.
87106
writeSize:=int(bps/10)
88107
rch:=make(chanerror)
89108
wch:=make(chanerror)
109+
110+
// Read forever in the background.
90111
gofunc() {
91-
rch<-readForever(ctx,&crw)
112+
rch<-readContext(deadlineCtx,&crw,writeSize*2)
113+
conn.Close()
92114
close(rch)
93115
}()
116+
117+
// Write random data to the PTY every tick.
94118
gofunc() {
95-
wch<-writeRandomData(ctx,&crw,writeSize,100*time.Millisecond)
119+
wch<-writeRandomData(deadlineCtx,&crw,writeSize,tick.C)
96120
close(wch)
97121
}()
98122

123+
// Wait for both our reads and writes to be finished.
99124
ifwErr:=<-wch;wErr!=nil {
100125
returnxerrors.Errorf("write to pty: %w",wErr)
101126
}
102127
ifrErr:=<-rch;rErr!=nil {
103128
returnxerrors.Errorf("read from pty: %w",rErr)
104129
}
105130

106-
_,_=fmt.Fprintf(inv.Stdout,"Test results:\n")
107-
_,_=fmt.Fprintf(inv.Stdout,"Took: %.2fs\n",time.Since(start).Seconds())
108-
_,_=fmt.Fprintf(inv.Stdout,"Sent: %d bytes\n",crw.BytesWritten())
109-
_,_=fmt.Fprintf(inv.Stdout,"Rcvd: %d bytes\n",crw.BytesRead())
110-
returnnil
131+
duration:=time.Since(start)
132+
133+
results:=trafficGenOutput{
134+
DurationSeconds:duration.Seconds(),
135+
SentBytes:crw.BytesWritten(),
136+
RcvdBytes:crw.BytesRead(),
137+
}
138+
139+
out,err:=formatter.Format(inv.Context(),results)
140+
iferr!=nil {
141+
returnerr
142+
}
143+
144+
_,err=fmt.Fprintln(inv.Stdout,out)
145+
returnerr
111146
},
112147
}
113148

@@ -128,66 +163,78 @@ func (r *RootCmd) trafficGen() *clibase.Cmd {
128163
},
129164
}
130165

166+
formatter.AttachOptions(&cmd.Options)
131167
returncmd
132168
}
133169

134-
funcreadForever(ctx context.Context,src io.Reader)error {
135-
buf:=make([]byte,1024)
170+
funcreadContext(ctx context.Context,src io.Reader,bufSizeint)error {
171+
buf:=make([]byte,bufSize)
136172
for {
137173
select {
138174
case<-ctx.Done():
139175
returnnil
140176
default:
177+
ifctx.Err()!=nil {
178+
returnnil
179+
}
141180
_,err:=src.Read(buf)
142-
iferr!=nil&&err!=io.EOF {
181+
iferr!=nil {
182+
ifxerrors.Is(err,io.EOF) {
183+
returnnil
184+
}
143185
returnerr
144186
}
145187
}
146188
}
147189
}
148190

149-
funcwriteRandomData(ctx context.Context,dst io.Writer,sizeint,period time.Duration)error {
150-
tick:=time.NewTicker(period)
151-
defertick.Stop()
191+
funcwriteRandomData(ctx context.Context,dst io.Writer,sizeint,tick<-chan time.Time)error {
152192
for {
153193
select {
154194
case<-ctx.Done():
155195
returnnil
156-
case<-tick.C:
157-
randStr,err:=cryptorand.String(size)
158-
iferr!=nil {
159-
returnerr
160-
}
196+
case<-tick:
197+
payload:="#"+mustRandStr(size-1)
161198
data,err:=json.Marshal(codersdk.ReconnectingPTYRequest{
162-
Data:randStr,
199+
Data:payload,
163200
})
164201
iferr!=nil {
165202
returnerr
166203
}
167-
err=copyContext(ctx,dst,data)
168-
iferr!=nil {
204+
if_,err:=copyContext(ctx,dst,data);err!=nil {
169205
returnerr
170206
}
171207
}
172208
}
173209
}
174210

175-
funccopyContext(ctx context.Context,dst io.Writer,src []byte)error {
176-
foridx:=rangesrc {
211+
// copyContext copies from src to dst until ctx is canceled.
212+
funccopyContext(ctx context.Context,dst io.Writer,src []byte) (int,error) {
213+
varcountint
214+
for {
177215
select {
178216
case<-ctx.Done():
179-
returnnil
217+
returncount,nil
180218
default:
181-
_,err:=dst.Write(src[idx :idx+1])
219+
ifctx.Err()!=nil {
220+
returncount,nil
221+
}
222+
n,err:=dst.Write(src)
182223
iferr!=nil {
183224
ifxerrors.Is(err,io.EOF) {
184-
returnnil
225+
// On an EOF, assume that all of src was consumed.
226+
returnlen(src),nil
185227
}
186-
returnerr
228+
returncount,err
229+
}
230+
count+=n
231+
ifn==len(src) {
232+
returncount,nil
187233
}
234+
// Not all of src was consumed. Update src and retry.
235+
src=src[n:]
188236
}
189237
}
190-
returnnil
191238
}
192239

193240
typecountReadWriterstruct {
@@ -219,3 +266,11 @@ func (w *countReadWriter) BytesRead() int64 {
219266
func (w*countReadWriter)BytesWritten()int64 {
220267
returnw.bytesWritten.Load()
221268
}
269+
270+
funcmustRandStr(lenint)string {
271+
randStr,err:=cryptorand.String(len)
272+
iferr!=nil {
273+
panic(err)
274+
}
275+
returnrandStr
276+
}

‎cli/trafficgen_test.go

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ package cli_test
33
import (
44
"bytes"
55
"context"
6-
"strings"
6+
"encoding/json"
77
"testing"
88

99
"github.com/google/uuid"
@@ -23,7 +23,6 @@ import (
2323
// We do not perform any cleanup.
2424
funcTestTrafficGen(t*testing.T) {
2525
t.Parallel()
26-
t.Skip("TODO: this hangs in a unit test but works in the real world.")
2726

2827
ctx,cancelFunc:=context.WithTimeout(context.Background(),testutil.WaitMedium)
2928
defercancelFunc()
@@ -74,20 +73,22 @@ func TestTrafficGen(t *testing.T) {
7473
inv,root:=clitest.New(t,"trafficgen",ws.Name,
7574
"--duration","1s",
7675
"--bps","100",
76+
"-o","json",
7777
)
7878
clitest.SetupConfig(t,client,root)
7979
varstdout,stderr bytes.Buffer
8080
inv.Stdout=&stdout
8181
inv.Stderr=&stderr
8282
err:=inv.WithContext(ctx).Run()
8383
require.NoError(t,err)
84-
stdoutStr:=stdout.String()
85-
stderrStr:=stderr.String()
86-
require.Empty(t,stderrStr)
87-
lines:=strings.Split(strings.TrimSpace(stdoutStr),"\n")
88-
require.Len(t,lines,4)
89-
require.Equal(t,"Test results:",lines[0])
90-
require.Regexp(t,`Took:\s+\d+\.\d+s`,lines[1])
91-
require.Regexp(t,`Sent:\s+\d+ bytes`,lines[2])
92-
require.Regexp(t,`Rcvd:\s+\d+ bytes`,lines[3])
84+
// TODO: this struct is currently unexported. Put it somewhere better.
85+
varoutputstruct {
86+
DurationSecondsfloat64`json:"duration_s"`
87+
SentBytesint64`json:"sent_bytes"`
88+
RcvdBytesint64`json:"rcvd_bytes"`
89+
}
90+
require.NoError(t,json.Unmarshal(stdout.Bytes(),&output))
91+
require.NotZero(t,output.DurationSeconds)
92+
require.NotZero(t,output.SentBytes)
93+
require.NotZero(t,output.RcvdBytes)
9394
}

0 commit comments

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