feat: Add speedtest command for tailnet (#3874)

cli/root.go

@@ -78,6 +78,7 @@ func Core() []*cobra.Command {
 		schedules(),
 		show(),
 		ssh(),
+		speedtest(),
 		start(),
 		state(),
 		stop(),

cli/speedtest.go

@@ -0,0 +1,91 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/cli/cliflag"
+	"github.com/coder/coder/cli/cliui"
+	"github.com/coder/coder/codersdk"
+	"github.com/jedib0t/go-pretty/v6/table"
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+	tsspeedtest "tailscale.com/net/speedtest"
+)
+
+func speedtest() *cobra.Command {
+	var (
+		reverse bool
+		timeStr string
+	)
+	cmd := &cobra.Command{
+		Annotations: workspaceCommand,
+		Use:         "speedtest <workspace>",
+		Short:       "Run a speed test from your machine to the workspace.",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx, cancel := context.WithCancel(cmd.Context())
+			defer cancel()
+
+			dur, err := time.ParseDuration(timeStr)
+			if err != nil {
+				return err
+			}
+
+			client, err := CreateClient(cmd)
+			if err != nil {
+				return xerrors.Errorf("create codersdk client: %w", err)
+			}
+
+			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, cmd, client, codersdk.Me, args[0], false)
+			if err != nil {
+				return err
+			}
+
+			err = cliui.Agent(ctx, cmd.ErrOrStderr(), cliui.AgentOptions{
+				WorkspaceName: workspace.Name,
+				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
+					return client.WorkspaceAgent(ctx, workspaceAgent.ID)
+				},
+			})
+			if err != nil {
+				return xerrors.Errorf("await agent: %w", err)
+			}
+			conn, err := client.DialWorkspaceAgentTailnet(ctx, slog.Logger{}, workspaceAgent.ID)
+			if err != nil {
+				return err
+			}
+			defer conn.Close()
+			_, _ = conn.Ping()
+			dir := tsspeedtest.Download
+			if reverse {
+				dir = tsspeedtest.Upload
+			}
+			cmd.Printf("Starting a %ds %s test...\n", int(dur.Seconds()), dir)
+			results, err := conn.Speedtest(dir, dur)
+			if err != nil {
+				return err
+			}
+			tableWriter := cliui.Table()
+			tableWriter.AppendHeader(table.Row{"Interval", "Transfer", "Bandwidth"})
+			for _, r := range results {
+				if r.Total {
+					tableWriter.AppendSeparator()
+				}
+				tableWriter.AppendRow(table.Row{
+					fmt.Sprintf("%.2f-%.2f sec", r.IntervalStart.Seconds(), r.IntervalEnd.Seconds()),
+					fmt.Sprintf("%.4f MBits", r.MegaBits()),
+					fmt.Sprintf("%.4f Mbits/sec", r.MBitsPerSecond()),
+				})
+			}
+			_, err = fmt.Fprintln(cmd.OutOrStdout(), tableWriter.Render())
+			return err
+		},
+	}
+	cliflag.BoolVarP(cmd.Flags(), &reverse, "reverse", "r", "", false,
+		"Specifies whether to run in reverse mode where the client receives and the server sends.")
+	cliflag.StringVarP(cmd.Flags(), &timeStr, "time", "t", "", "5s",
+		"Specifies the duration to monitor traffic.")
+	return cmd
+}

cli/speedtest_test.go

@@ -0,0 +1,46 @@
+package cli_test
+
+import (
+	"context"
+	"testing"
+
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/agent"
+	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/pty/ptytest"
+	"github.com/coder/coder/testutil"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSpeedtest(t *testing.T) {
+	t.Parallel()
+	if testing.Short() {
+		t.Skip("This test takes a minimum of 5ms per a hardcoded value in Tailscale!")
+	}
+	client, workspace, agentToken := setupWorkspaceForAgent(t)
+	agentClient := codersdk.New(client.URL)
+	agentClient.SessionToken = agentToken
+	agentCloser := agent.New(agent.Options{
+		FetchMetadata:     agentClient.WorkspaceAgentMetadata,
+		WebRTCDialer:      agentClient.ListenWorkspaceAgent,
+		CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
+		Logger:            slogtest.Make(t, nil).Named("agent"),
+	})
+	defer agentCloser.Close()
+	coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
+
+	cmd, root := clitest.New(t, "speedtest", workspace.Name)
+	clitest.SetupConfig(t, client, root)
+	pty := ptytest.New(t)
+	cmd.SetOut(pty.Output())
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+	defer cancel()
+	cmdDone := tGo(t, func() {
+		err := cmd.ExecuteContext(ctx)
+		assert.NoError(t, err)
+	})
+	<-cmdDone
+}

cli/ssh_test.go

@@ -31,7 +31,7 @@ import (
 	"github.com/coder/coder/testutil"
 )
 
-func setupWorkspaceForSSH(t *testing.T) (*codersdk.Client, codersdk.Workspace, string) {
+func setupWorkspaceForAgent(t *testing.T) (*codersdk.Client, codersdk.Workspace, string) {
 	t.Helper()
 	client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
 	user := coderdtest.CreateFirstUser(t, client)
@@ -69,7 +69,7 @@ func TestSSH(t *testing.T) {
 	t.Run("ImmediateExit", func(t *testing.T) {
 		t.Parallel()
 
-		client, workspace, agentToken := setupWorkspaceForSSH(t)
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
 		cmd, root := clitest.New(t, "ssh", workspace.Name)
 		clitest.SetupConfig(t, client, root)
 		pty := ptytest.New(t)
@@ -104,7 +104,7 @@ func TestSSH(t *testing.T) {
 	})
 	t.Run("Stdio", func(t *testing.T) {
 		t.Parallel()
-		client, workspace, agentToken := setupWorkspaceForSSH(t)
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
 		_, _ = tGoContext(t, func(ctx context.Context) {
 			// Run this async so the SSH command has to wait for
 			// the build and agent to connect!
@@ -175,7 +175,7 @@ func TestSSH(t *testing.T) {
 
 		t.Parallel()
 
-		client, workspace, agentToken := setupWorkspaceForSSH(t)
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
 
 		agentClient := codersdk.New(client.URL)
 		agentClient.SessionToken = agentToken