Skip to content

Commit 5f18499

Browse files
deansheatherpull[bot]
deansheather
authored and
pull[bot]
committed
feat: add reconnectingpty loadtest (#5083)
1 parent ac521f0 commit 5f18499

File tree

11 files changed

+607
-20
lines changed

11 files changed

+607
-20
lines changed

‎agent/agent.go

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -725,6 +725,7 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
725725
func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.ReconnectingPTYInit, conn net.Conn) {
726726
defer conn.Close()
727727

728+
connectionID := uuid.NewString()
728729
var rpty *reconnectingPTY
729730
rawRPTY, ok := a.reconnectingPTYs.Load(msg.ID)
730731
if ok {
@@ -760,8 +761,12 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
760761
a.closeMutex.Unlock()
761762
ctx, cancelFunc := context.WithCancel(ctx)
762763
rpty = &reconnectingPTY{
763-
activeConns: make(map[string]net.Conn),
764-
ptty: ptty,
764+
activeConns: map[string]net.Conn{
765+
// We have to put the connection in the map instantly otherwise
766+
// the connection won't be closed if the process instantly dies.
767+
connectionID: conn,
768+
},
769+
ptty: ptty,
765770
// Timeouts created with an after func can be reset!
766771
timeout: time.AfterFunc(a.reconnectingPTYTimeout, cancelFunc),
767772
circularBuffer: circularBuffer,
@@ -827,7 +832,6 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
827832
a.logger.Warn(ctx, "write reconnecting pty buffer", slog.F("id", msg.ID), slog.Error(err))
828833
return
829834
}
830-
connectionID := uuid.NewString()
831835
// Multiple connections to the same TTY are permitted.
832836
// This could easily be used for terminal sharing, but
833837
// we do it because it's a nice user experience to

‎agent/agent_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ func TestAgent(t *testing.T) {
8383

8484
conn, stats, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
8585

86-
ptyConn, err := conn.ReconnectingPTY(ctx, uuid.NewString(), 128, 128, "/bin/bash")
86+
ptyConn, err := conn.ReconnectingPTY(ctx, uuid.New(), 128, 128, "/bin/bash")
8787
require.NoError(t, err)
8888
defer ptyConn.Close()
8989

@@ -405,7 +405,7 @@ func TestAgent(t *testing.T) {
405405
defer cancel()
406406

407407
conn, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
408-
id := uuid.NewString()
408+
id := uuid.New()
409409
netConn, err := conn.ReconnectingPTY(ctx, id, 100, 100, "/bin/bash")
410410
require.NoError(t, err)
411411
bufRead := bufio.NewReader(netConn)

‎cli/loadtestconfig.go

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"github.com/coder/coder/loadtest/agentconn"
1111
"github.com/coder/coder/loadtest/harness"
1212
"github.com/coder/coder/loadtest/placebo"
13+
"github.com/coder/coder/loadtest/reconnectingpty"
1314
"github.com/coder/coder/loadtest/workspacebuild"
1415
)
1516

@@ -88,9 +89,10 @@ func (s LoadTestStrategy) ExecutionStrategy() harness.ExecutionStrategy {
8889
type LoadTestType string
8990

9091
const (
91-
LoadTestTypeAgentConn LoadTestType = "agentconn"
92-
LoadTestTypePlacebo LoadTestType = "placebo"
93-
LoadTestTypeWorkspaceBuild LoadTestType = "workspacebuild"
92+
LoadTestTypeAgentConn LoadTestType = "agentconn"
93+
LoadTestTypePlacebo LoadTestType = "placebo"
94+
LoadTestTypeReconnectingPTY LoadTestType = "reconnectingpty"
95+
LoadTestTypeWorkspaceBuild LoadTestType = "workspacebuild"
9496
)
9597

9698
type LoadTest struct {
@@ -104,6 +106,8 @@ type LoadTest struct {
104106
AgentConn *agentconn.Config `json:"agentconn,omitempty"`
105107
// Placebo must be set if type == "placebo".
106108
Placebo *placebo.Config `json:"placebo,omitempty"`
109+
// ReconnectingPTY must be set if type == "reconnectingpty".
110+
ReconnectingPTY *reconnectingpty.Config `json:"reconnectingpty,omitempty"`
107111
// WorkspaceBuild must be set if type == "workspacebuild".
108112
WorkspaceBuild *workspacebuild.Config `json:"workspacebuild,omitempty"`
109113
}
@@ -120,6 +124,11 @@ func (t LoadTest) NewRunner(client *codersdk.Client) (harness.Runnable, error) {
120124
return nil, xerrors.New("placebo config must be set")
121125
}
122126
return placebo.NewRunner(*t.Placebo), nil
127+
case LoadTestTypeReconnectingPTY:
128+
if t.ReconnectingPTY == nil {
129+
return nil, xerrors.New("reconnectingpty config must be set")
130+
}
131+
return reconnectingpty.NewRunner(client, *t.ReconnectingPTY), nil
123132
case LoadTestTypeWorkspaceBuild:
124133
if t.WorkspaceBuild == nil {
125134
return nil, xerrors.Errorf("workspacebuild config must be set")
@@ -185,6 +194,15 @@ func (t *LoadTest) Validate() error {
185194
if err != nil {
186195
return xerrors.Errorf("validate placebo: %w", err)
187196
}
197+
case LoadTestTypeReconnectingPTY:
198+
if t.ReconnectingPTY == nil {
199+
return xerrors.Errorf("reconnectingpty test type must specify reconnectingpty")
200+
}
201+
202+
err := t.ReconnectingPTY.Validate()
203+
if err != nil {
204+
return xerrors.Errorf("validate reconnectingpty: %w", err)
205+
}
188206
case LoadTestTypeWorkspaceBuild:
189207
if t.WorkspaceBuild == nil {
190208
return xerrors.New("workspacebuild test type must specify workspacebuild")

‎coderd/workspaceagents.go

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import (
66
"encoding/json"
77
"errors"
88
"fmt"
9-
"io"
109
"net"
1110
"net/http"
1211
"net/netip"
@@ -26,6 +25,7 @@ import (
2625
"tailscale.com/tailcfg"
2726

2827
"cdr.dev/slog"
28+
"github.com/coder/coder/agent"
2929
"github.com/coder/coder/coderd/database"
3030
"github.com/coder/coder/coderd/gitauth"
3131
"github.com/coder/coder/coderd/httpapi"
@@ -247,17 +247,13 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
247247
return
248248
}
249249
defer release()
250-
ptNetConn, err := agentConn.ReconnectingPTY(ctx, reconnect.String(), uint16(height), uint16(width), r.URL.Query().Get("command"))
250+
ptNetConn, err := agentConn.ReconnectingPTY(ctx, reconnect, uint16(height), uint16(width), r.URL.Query().Get("command"))
251251
if err != nil {
252252
_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial: %s", err))
253253
return
254254
}
255255
defer ptNetConn.Close()
256-
// Pipe the ends together!
257-
go func() {
258-
_, _ = io.Copy(wsNetConn, ptNetConn)
259-
}()
260-
_, _ = io.Copy(ptNetConn, wsNetConn)
256+
agent.Bicopy(ctx, wsNetConn, ptNetConn)
261257
}
262258

263259
func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Request) {

‎codersdk/agentconn.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"strings"
1515
"time"
1616

17+
"github.com/google/uuid"
1718
"golang.org/x/crypto/ssh"
1819
"golang.org/x/xerrors"
1920
"tailscale.com/net/speedtest"
@@ -158,13 +159,13 @@ func (c *AgentConn) Close() error {
158159

159160
// @typescript-ignore ReconnectingPTYInit
160161
type ReconnectingPTYInit struct {
161-
ID string
162+
ID uuid.UUID
162163
Height uint16
163164
Width uint16
164165
Command string
165166
}
166167

167-
func (c *AgentConn) ReconnectingPTY(ctx context.Context, id string, height, width uint16, command string) (net.Conn, error) {
168+
func (c *AgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, width uint16, command string) (net.Conn, error) {
168169
ctx, span := tracing.StartSpan(ctx)
169170
defer span.End()
170171

‎codersdk/client.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -268,7 +268,7 @@ func readBodyAsError(res *http.Response) error {
268268
return &Error{
269269
statusCode: res.StatusCode,
270270
Response: Response{
271-
Message: "unexpected non-JSON response",
271+
Message: fmt.Sprintf("unexpected non-JSON response %q", contentType),
272272
Detail: string(resp),
273273
},
274274
Helper: helper,

‎codersdk/workspaceagents.go

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -501,11 +501,18 @@ func (c *Client) PostWorkspaceAgentVersion(ctx context.Context, version string)
501501
// WorkspaceAgentReconnectingPTY spawns a PTY that reconnects using the token provided.
502502
// It communicates using `agent.ReconnectingPTYRequest` marshaled as JSON.
503503
// Responses are PTY output that can be rendered.
504-
func (c *Client) WorkspaceAgentReconnectingPTY(ctx context.Context, agentID, reconnect uuid.UUID, height, width int, command string) (net.Conn, error) {
505-
serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/pty?reconnect=%s&height=%d&width=%d&command=%s", agentID, reconnect, height, width, command))
504+
func (c *Client) WorkspaceAgentReconnectingPTY(ctx context.Context, agentID, reconnect uuid.UUID, height, width uint16, command string) (net.Conn, error) {
505+
serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/pty", agentID))
506506
if err != nil {
507507
return nil, xerrors.Errorf("parse url: %w", err)
508508
}
509+
q := serverURL.Query()
510+
q.Set("reconnect", reconnect.String())
511+
q.Set("height", strconv.Itoa(int(height)))
512+
q.Set("width", strconv.Itoa(int(width)))
513+
q.Set("command", command)
514+
serverURL.RawQuery = q.Encode()
515+
509516
jar, err := cookiejar.New(nil)
510517
if err != nil {
511518
return nil, xerrors.Errorf("create cookie jar: %w", err)

‎loadtest/reconnectingpty/config.go

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
package reconnectingpty
2+
3+
import (
4+
"time"
5+
6+
"github.com/google/uuid"
7+
"golang.org/x/xerrors"
8+
9+
"github.com/coder/coder/coderd/httpapi"
10+
"github.com/coder/coder/codersdk"
11+
)
12+
13+
const (
14+
DefaultWidth = 80
15+
DefaultHeight = 24
16+
DefaultTimeout = httpapi.Duration(5 * time.Minute)
17+
)
18+
19+
type Config struct {
20+
// AgentID is the ID of the agent to run the command in.
21+
AgentID uuid.UUID `json:"agent_id"`
22+
// Init is the initial packet to send to the agent when launching the TTY.
23+
// If the ID is not set, defaults to a random UUID. If the width or height
24+
// is not set, defaults to 80x24. If the command is not set, defaults to
25+
// opening a login shell. Command runs in the default shell.
26+
Init codersdk.ReconnectingPTYInit `json:"init"`
27+
// Timeout is the duration to wait for the command to exit. Defaults to
28+
// 5 minutes.
29+
Timeout httpapi.Duration `json:"timeout"`
30+
// ExpectTimeout means we expect the timeout to be reached (i.e. the command
31+
// doesn't exit within the given timeout).
32+
ExpectTimeout bool `json:"expect_timeout"`
33+
// ExpectOutput checks that the given string is present in the output. The
34+
// string must be present on a single line.
35+
ExpectOutput string `json:"expect_output"`
36+
// LogOutput determines whether the output of the command should be logged.
37+
// For commands that produce a lot of output this should be disabled to
38+
// avoid loadtest OOMs. All log output is still read and discarded if this
39+
// is false.
40+
LogOutput bool `json:"log_output"`
41+
}
42+
43+
func (c Config) Validate() error {
44+
if c.AgentID == uuid.Nil {
45+
return xerrors.New("agent_id must be set")
46+
}
47+
if c.Timeout < 0 {
48+
return xerrors.New("timeout must be a positive value")
49+
}
50+
51+
return nil
52+
}

‎loadtest/reconnectingpty/config_test.go

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package reconnectingpty_test
2+
3+
import (
4+
"testing"
5+
"time"
6+
7+
"github.com/google/uuid"
8+
"github.com/stretchr/testify/require"
9+
10+
"github.com/coder/coder/coderd/httpapi"
11+
"github.com/coder/coder/codersdk"
12+
"github.com/coder/coder/loadtest/reconnectingpty"
13+
)
14+
15+
func Test_Config(t *testing.T) {
16+
t.Parallel()
17+
18+
id := uuid.New()
19+
cases := []struct {
20+
name string
21+
config reconnectingpty.Config
22+
errContains string
23+
}{
24+
{
25+
name: "OKBasic",
26+
config: reconnectingpty.Config{
27+
AgentID: id,
28+
},
29+
},
30+
{
31+
name: "OKFull",
32+
config: reconnectingpty.Config{
33+
AgentID: id,
34+
Init: codersdk.ReconnectingPTYInit{
35+
ID: id,
36+
Width: 80,
37+
Height: 24,
38+
Command: "echo 'hello world'",
39+
},
40+
Timeout: httpapi.Duration(time.Minute),
41+
ExpectTimeout: false,
42+
ExpectOutput: "hello world",
43+
LogOutput: true,
44+
},
45+
},
46+
{
47+
name: "NoAgentID",
48+
config: reconnectingpty.Config{
49+
AgentID: uuid.Nil,
50+
},
51+
errContains: "agent_id must be set",
52+
},
53+
{
54+
name: "NegativeTimeout",
55+
config: reconnectingpty.Config{
56+
AgentID: id,
57+
Timeout: httpapi.Duration(-time.Minute),
58+
},
59+
errContains: "timeout must be a positive value",
60+
},
61+
}
62+
63+
for _, c := range cases {
64+
c := c
65+
66+
t.Run(c.name, func(t *testing.T) {
67+
t.Parallel()
68+
69+
err := c.config.Validate()
70+
if c.errContains != "" {
71+
require.Error(t, err)
72+
require.Contains(t, err.Error(), c.errContains)
73+
} else {
74+
require.NoError(t, err)
75+
}
76+
})
77+
}
78+
}

0 commit comments

Comments
 (0)