feat(scaletest): add option to send traffic over SSH (#8521)

johnstcn · web-flow · commit 278527cff497 · 2023-07-18T12:17:11.000+01:00
- Refactors the metrics logic to avoid needing to pass in a whole prometheus registry - Adds an --ssh option to the workspace-traffic command to send SSH traffic Fixes #8242
diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go
@@ -848,6 +848,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
 	var (
 		tickInterval               time.Duration
 		bytesPerTick               int64
+		ssh                        bool
 		scaletestPrometheusAddress string
 		scaletestPrometheusWait    time.Duration
 
@@ -938,20 +939,19 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
 
 				// Setup our workspace agent connection.
 				config := workspacetraffic.Config{
-					AgentID:        agentID,
-					AgentName:      agentName,
-					BytesPerTick:   bytesPerTick,
-					Duration:       strategy.timeout,
-					TickInterval:   tickInterval,
-					WorkspaceName:  ws.Name,
-					WorkspaceOwner: ws.OwnerName,
-					Registry:       reg,
+					AgentID:      agentID,
+					BytesPerTick: bytesPerTick,
+					Duration:     strategy.timeout,
+					TickInterval: tickInterval,
+					ReadMetrics:  metrics.ReadMetrics(ws.OwnerName, ws.Name, agentName),
+					WriteMetrics: metrics.WriteMetrics(ws.OwnerName, ws.Name, agentName),
+					SSH:          ssh,
 				}
 
 				if err := config.Validate(); err != nil {
 					return xerrors.Errorf("validate config: %w", err)
 				}
-				var runner harness.Runnable = workspacetraffic.NewRunner(client, config, metrics)
+				var runner harness.Runnable = workspacetraffic.NewRunner(client, config)
 				if tracingEnabled {
 					runner = &runnableTraceWrapper{
 						tracer:   tracer,
@@ -1002,6 +1002,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd {
 			Description: "How often to send traffic.",
 			Value:       clibase.DurationOf(&tickInterval),
 		},
+		{
+			Flag:        "ssh",
+			Env:         "CODER_SCALETEST_WORKSPACE_TRAFFIC_SSH",
+			Default:     "",
+			Description: "Send traffic over SSH.",
+			Value:       clibase.BoolOf(&ssh),
+		},
 		{
 			Flag:        "scaletest-prometheus-address",
 			Env:         "CODER_SCALETEST_PROMETHEUS_ADDRESS",
diff --git a/cli/exp_scaletest_test.go b/cli/exp_scaletest_test.go
@@ -69,6 +69,7 @@ func TestScaleTestWorkspaceTraffic(t *testing.T) {
 		"--tick-interval", "100ms",
 		"--scaletest-prometheus-address", "127.0.0.1:0",
 		"--scaletest-prometheus-wait", "0s",
+		"--ssh",
 	)
 	clitest.SetupConfig(t, client, root)
 	var stdout, stderr bytes.Buffer
diff --git a/scaletest/workspacetraffic/config.go b/scaletest/workspacetraffic/config.go
@@ -4,23 +4,13 @@ import (
 	"time"
 
 	"github.com/google/uuid"
-	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/xerrors"
 )
 
 type Config struct {
 	// AgentID is the workspace agent ID to which to connect.
 	AgentID uuid.UUID `json:"agent_id"`
 
-	// AgentName is the name of the agent. Used for metrics.
-	AgentName string `json:"agent_name"`
-
-	// WorkspaceName is the name of the workspace. Used for metrics.
-	WorkspaceName string `json:"workspace_name"`
-
-	// WorkspaceOwner is the owner of the workspace. Used for metrics.
-	WorkspaceOwner string `json:"workspace_owner"`
-
 	// BytesPerTick is the number of bytes to send to the agent per tick.
 	BytesPerTick int64 `json:"bytes_per_tick"`
 
@@ -31,8 +21,10 @@ type Config struct {
 	// send data to workspace agents).
 	TickInterval time.Duration `json:"tick_interval"`
 
-	// Registry is a prometheus.Registerer for logging metrics
-	Registry prometheus.Registerer
+	ReadMetrics  ConnMetrics `json:"-"`
+	WriteMetrics ConnMetrics `json:"-"`
+
+	SSH bool `json:"ssh"`
 }
 
 func (c Config) Validate() error {
diff --git a/scaletest/workspacetraffic/conn.go b/scaletest/workspacetraffic/conn.go
@@ -0,0 +1,123 @@
+package workspacetraffic
+
+import (
+	"context"
+	"io"
+	"sync"
+
+	"github.com/coder/coder/codersdk"
+
+	"github.com/google/uuid"
+	"github.com/hashicorp/go-multierror"
+	gossh "golang.org/x/crypto/ssh"
+	"golang.org/x/xerrors"
+)
+
+func connectPTY(ctx context.Context, client *codersdk.Client, agentID, reconnect uuid.UUID) (*countReadWriteCloser, error) {
+	conn, err := client.WorkspaceAgentReconnectingPTY(ctx, codersdk.WorkspaceAgentReconnectingPTYOpts{
+		AgentID:   agentID,
+		Reconnect: reconnect,
+		Height:    25,
+		Width:     80,
+		Command:   "/bin/sh",
+	})
+	if err != nil {
+		return nil, xerrors.Errorf("connect pty: %w", err)
+	}
+
+	// Wrap the conn in a countReadWriteCloser so we can monitor bytes sent/rcvd.
+	crw := countReadWriteCloser{ctx: ctx, rwc: conn}
+	return &crw, nil
+}
+
+func connectSSH(ctx context.Context, client *codersdk.Client, agentID uuid.UUID) (*countReadWriteCloser, error) {
+	agentConn, err := client.DialWorkspaceAgent(ctx, agentID, &codersdk.DialWorkspaceAgentOptions{})
+	if err != nil {
+		return nil, xerrors.Errorf("dial workspace agent: %w", err)
+	}
+	agentConn.AwaitReachable(ctx)
+	sshClient, err := agentConn.SSHClient(ctx)
+	if err != nil {
+		return nil, xerrors.Errorf("get ssh client: %w", err)
+	}
+	sshSession, err := sshClient.NewSession()
+	if err != nil {
+		_ = agentConn.Close()
+		return nil, xerrors.Errorf("new ssh session: %w", err)
+	}
+	wrappedConn := &wrappedSSHConn{ctx: ctx}
+	// Do some plumbing to hook up the wrappedConn
+	pr1, pw1 := io.Pipe()
+	wrappedConn.stdout = pr1
+	sshSession.Stdout = pw1
+	pr2, pw2 := io.Pipe()
+	sshSession.Stdin = pr2
+	wrappedConn.stdin = pw2
+	err = sshSession.RequestPty("xterm", 25, 80, gossh.TerminalModes{})
+	if err != nil {
+		_ = pr1.Close()
+		_ = pr2.Close()
+		_ = pw1.Close()
+		_ = pw2.Close()
+		_ = sshSession.Close()
+		_ = agentConn.Close()
+		return nil, xerrors.Errorf("request pty: %w", err)
+	}
+	err = sshSession.Shell()
+	if err != nil {
+		_ = sshSession.Close()
+		_ = agentConn.Close()
+		return nil, xerrors.Errorf("shell: %w", err)
+	}
+
+	closeFn := func() error {
+		var merr error
+		if err := sshSession.Close(); err != nil {
+			merr = multierror.Append(merr, err)
+		}
+		if err := agentConn.Close(); err != nil {
+			merr = multierror.Append(merr, err)
+		}
+		return merr
+	}
+	wrappedConn.close = closeFn
+
+	crw := &countReadWriteCloser{ctx: ctx, rwc: wrappedConn}
+	return crw, nil
+}
+
+// wrappedSSHConn wraps an ssh.Session to implement io.ReadWriteCloser.
+type wrappedSSHConn struct {
+	ctx       context.Context
+	stdout    io.Reader
+	stdin     io.Writer
+	closeOnce sync.Once
+	closeErr  error
+	close     func() error
+}
+
+func (w *wrappedSSHConn) Close() error {
+	w.closeOnce.Do(func() {
+		_, _ = w.stdin.Write([]byte("exit\n"))
+		w.closeErr = w.close()
+	})
+	return w.closeErr
+}
+
+func (w *wrappedSSHConn) Read(p []byte) (n int, err error) {
+	select {
+	case <-w.ctx.Done():
+		return 0, xerrors.Errorf("read: %w", w.ctx.Err())
+	default:
+		return w.stdout.Read(p)
+	}
+}
+
+func (w *wrappedSSHConn) Write(p []byte) (n int, err error) {
+	select {
+	case <-w.ctx.Done():
+		return 0, xerrors.Errorf("write: %w", w.ctx.Err())
+	default:
+		return w.stdin.Write(p)
+	}
+}
diff --git a/scaletest/workspacetraffic/countreadwriter.go b/scaletest/workspacetraffic/countreadwriter.go
@@ -0,0 +1,71 @@
+package workspacetraffic
+
+import (
+	"context"
+	"errors"
+	"io"
+	"time"
+
+	"golang.org/x/xerrors"
+
+	"nhooyr.io/websocket"
+)
+
+// countReadWriteCloser wraps an io.ReadWriteCloser and counts the number of bytes read and written.
+type countReadWriteCloser struct {
+	ctx          context.Context
+	rwc          io.ReadWriteCloser
+	readMetrics  ConnMetrics
+	writeMetrics ConnMetrics
+}
+
+func (w *countReadWriteCloser) Close() error {
+	return w.rwc.Close()
+}
+
+func (w *countReadWriteCloser) Read(p []byte) (int, error) {
+	start := time.Now()
+	n, err := w.rwc.Read(p)
+	took := time.Since(start).Seconds()
+	if reportableErr(err) {
+		w.readMetrics.AddError(1)
+	}
+	w.readMetrics.ObserveLatency(took)
+	if n > 0 {
+		w.readMetrics.AddTotal(float64(n))
+	}
+	return n, err
+}
+
+func (w *countReadWriteCloser) Write(p []byte) (int, error) {
+	start := time.Now()
+	n, err := w.rwc.Write(p)
+	took := time.Since(start).Seconds()
+	if reportableErr(err) {
+		w.writeMetrics.AddError(1)
+	}
+	w.writeMetrics.ObserveLatency(took)
+	if n > 0 {
+		w.writeMetrics.AddTotal(float64(n))
+	}
+	return n, err
+}
+
+// some errors we want to report in metrics; others we want to ignore
+// such as websocket.StatusNormalClosure or context.Canceled
+func reportableErr(err error) bool {
+	if err == nil {
+		return false
+	}
+	if xerrors.Is(err, io.EOF) {
+		return false
+	}
+	if xerrors.Is(err, context.Canceled) {
+		return false
+	}
+	var wsErr websocket.CloseError
+	if errors.As(err, &wsErr) {
+		return wsErr.Code != websocket.StatusNormalClosure
+	}
+	return false
+}
diff --git a/scaletest/workspacetraffic/metrics.go b/scaletest/workspacetraffic/metrics.go
@@ -54,3 +54,43 @@ func NewMetrics(reg prometheus.Registerer, labelNames ...string) *Metrics {
 	reg.MustRegister(m.WriteLatencySeconds)
 	return m
 }
+
+func (m *Metrics) ReadMetrics(lvs ...string) ConnMetrics {
+	return &connMetrics{
+		addError:       m.ReadErrorsTotal.WithLabelValues(lvs...).Add,
+		observeLatency: m.ReadLatencySeconds.WithLabelValues(lvs...).Observe,
+		addTotal:       m.BytesReadTotal.WithLabelValues(lvs...).Add,
+	}
+}
+
+func (m *Metrics) WriteMetrics(lvs ...string) ConnMetrics {
+	return &connMetrics{
+		addError:       m.WriteErrorsTotal.WithLabelValues(lvs...).Add,
+		observeLatency: m.WriteLatencySeconds.WithLabelValues(lvs...).Observe,
+		addTotal:       m.BytesWrittenTotal.WithLabelValues(lvs...).Add,
+	}
+}
+
+type ConnMetrics interface {
+	AddError(float64)
+	ObserveLatency(float64)
+	AddTotal(float64)
+}
+
+type connMetrics struct {
+	addError       func(float64)
+	observeLatency func(float64)
+	addTotal       func(float64)
+}
+
+func (c *connMetrics) AddError(f float64) {
+	c.addError(f)
+}
+
+func (c *connMetrics) ObserveLatency(f float64) {
+	c.observeLatency(f)
+}
+
+func (c *connMetrics) AddTotal(f float64) {
+	c.addTotal(f)
+}
diff --git a/scaletest/workspacetraffic/run.go b/scaletest/workspacetraffic/run.go
diff --git a/scaletest/workspacetraffic/run_test.go b/scaletest/workspacetraffic/run_test.go

Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,7 @@ func TestScaleTestWorkspaceTraffic(t *testing.T) {`
`69`	`69`	`"--tick-interval", "100ms",`
`70`	`70`	`"--scaletest-prometheus-address", "127.0.0.1:0",`
`71`	`71`	`"--scaletest-prometheus-wait", "0s",`
	`72`	`+ "--ssh",`
`72`	`73`	`)`
`73`	`74`	`clitest.SetupConfig(t, client, root)`
`74`	`75`	`var stdout, stderr bytes.Buffer`