From 1254e7a9026065183963cbf99be4a8284e28ba37 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Mon, 5 Sep 2022 17:15:49 -0500 Subject: [PATCH] feat: Add speedtest command for tailnet (#3874) --- agent/agent.go | 19 +++++++++ agent/agent_test.go | 16 ++++++++ agent/conn.go | 18 +++++++++ cli/root.go | 1 + cli/speedtest.go | 91 +++++++++++++++++++++++++++++++++++++++++++ cli/speedtest_test.go | 46 ++++++++++++++++++++++ cli/ssh_test.go | 8 ++-- go.mod | 4 +- go.sum | 4 +- 9 files changed, 199 insertions(+), 8 deletions(-) create mode 100644 cli/speedtest.go create mode 100644 cli/speedtest_test.go diff --git a/agent/agent.go b/agent/agent.go index 981c633a82..1274f4e5a7 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -29,6 +29,7 @@ import ( "go.uber.org/atomic" gossh "golang.org/x/crypto/ssh" "golang.org/x/xerrors" + "tailscale.com/net/speedtest" "tailscale.com/tailcfg" "cdr.dev/slog" @@ -58,6 +59,7 @@ var ( tailnetIP = netip.MustParseAddr("fd7a:115c:a1e0:49d6:b259:b7ac:b1b2:48f4") tailnetSSHPort = 1 tailnetReconnectingPTYPort = 2 + tailnetSpeedtestPort = 3 ) type Options struct { @@ -256,6 +258,23 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) { go a.handleReconnectingPTY(ctx, msg, conn) } }() + speedtestListener, err := a.network.Listen("tcp", ":"+strconv.Itoa(tailnetSpeedtestPort)) + if err != nil { + a.logger.Critical(ctx, "listen for speedtest", slog.Error(err)) + return + } + go func() { + for { + conn, err := speedtestListener.Accept() + if err != nil { + a.logger.Debug(ctx, "speedtest listener failed", slog.Error(err)) + return + } + go func() { + _ = speedtest.ServeConn(conn) + }() + } + }() } // runCoordinator listens for nodes and updates the self-node as it changes. diff --git a/agent/agent_test.go b/agent/agent_test.go index 49f57214ab..cac4c46330 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -19,6 +19,7 @@ import ( "time" "golang.org/x/xerrors" + "tailscale.com/net/speedtest" "tailscale.com/tailcfg" scp "github.com/bramvdbogaerde/go-scp" @@ -547,6 +548,21 @@ func TestAgent(t *testing.T) { return err == nil }, testutil.WaitMedium, testutil.IntervalFast) }) + + t.Run("Speedtest", func(t *testing.T) { + t.Parallel() + if testing.Short() { + t.Skip("The minimum duration for a speedtest is hardcoded in Tailscale to 5s!") + } + derpMap := tailnettest.RunDERPAndSTUN(t) + conn, _ := setupAgent(t, agent.Metadata{ + DERPMap: derpMap, + }, 0) + defer conn.Close() + res, err := conn.Speedtest(speedtest.Upload, speedtest.MinDuration) + require.NoError(t, err) + t.Logf("%.2f MBits/s", res[len(res)-1].MBitsPerSecond()) + }) } func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd { diff --git a/agent/conn.go b/agent/conn.go index 0e95e97e21..2e1b45ea0b 100644 --- a/agent/conn.go +++ b/agent/conn.go @@ -16,6 +16,7 @@ import ( "golang.org/x/crypto/ssh" "golang.org/x/xerrors" "tailscale.com/ipn/ipnstate" + "tailscale.com/net/speedtest" "tailscale.com/tailcfg" "github.com/coder/coder/peer" @@ -39,6 +40,7 @@ type Conn interface { CloseWithError(err error) error ReconnectingPTY(id string, height, width uint16, command string) (net.Conn, error) SSH() (net.Conn, error) + Speedtest(direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error) SSHClient() (*ssh.Client, error) DialContext(ctx context.Context, network string, addr string) (net.Conn, error) } @@ -77,6 +79,10 @@ func (c *WebRTCConn) SSH() (net.Conn, error) { return channel.NetConn(), nil } +func (*WebRTCConn) Speedtest(_ speedtest.Direction, _ time.Duration) ([]speedtest.Result, error) { + return nil, xerrors.New("not implemented") +} + // SSHClient calls SSH to create a client that uses a weak cipher // for high throughput. func (c *WebRTCConn) SSHClient() (*ssh.Client, error) { @@ -227,6 +233,18 @@ func (c *TailnetConn) SSHClient() (*ssh.Client, error) { return ssh.NewClient(sshConn, channels, requests), nil } +func (c *TailnetConn) Speedtest(direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error) { + speedConn, err := c.DialContextTCP(context.Background(), netip.AddrPortFrom(tailnetIP, uint16(tailnetSpeedtestPort))) + if err != nil { + return nil, xerrors.Errorf("dial speedtest: %w", err) + } + results, err := speedtest.RunClientWithConn(direction, duration, speedConn) + if err != nil { + return nil, xerrors.Errorf("run speedtest: %w", err) + } + return results, err +} + func (c *TailnetConn) DialContext(ctx context.Context, network string, addr string) (net.Conn, error) { _, rawPort, _ := net.SplitHostPort(addr) port, _ := strconv.Atoi(rawPort) diff --git a/cli/root.go b/cli/root.go index f90c3ad1ec..69575768ab 100644 --- a/cli/root.go +++ b/cli/root.go @@ -78,6 +78,7 @@ func Core() []*cobra.Command { schedules(), show(), ssh(), + speedtest(), start(), state(), stop(), diff --git a/cli/speedtest.go b/cli/speedtest.go new file mode 100644 index 0000000000..c9e69a2ddf --- /dev/null +++ b/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 ", + 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 +} diff --git a/cli/speedtest_test.go b/cli/speedtest_test.go new file mode 100644 index 0000000000..0eccd69171 --- /dev/null +++ b/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 +} diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 410c842970..be68ad6226 100644 --- a/cli/ssh_test.go +++ b/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 diff --git a/go.mod b/go.mod index 57d68cccaa..724fe34d52 100644 --- a/go.mod +++ b/go.mod @@ -49,7 +49,7 @@ replace github.com/tcnksm/go-httpstat => github.com/kylecarbs/go-httpstat v0.0.0 // There are a few minor changes we make to Tailscale that we're slowly upstreaming. Compare here: // https://github.com/tailscale/tailscale/compare/main...coder:tailscale:main -replace tailscale.com => github.com/coder/tailscale v1.1.1-0.20220902164407-ae46caa65076 +replace tailscale.com => github.com/coder/tailscale v1.1.1-0.20220905194158-291661887d25 require ( cdr.dev/slog v1.4.2-0.20220525200111-18dce5c2cd5f @@ -157,7 +157,7 @@ require ( k8s.io/utils v0.0.0-20220210201930-3a6ce19ff2f9 nhooyr.io/websocket v1.8.7 storj.io/drpc v0.0.33-0.20220622181519-9206537a4db7 - tailscale.com v1.26.2 + tailscale.com v1.30.0 ) require ( diff --git a/go.sum b/go.sum index d1915643f5..98b7b92c05 100644 --- a/go.sum +++ b/go.sum @@ -352,8 +352,8 @@ github.com/coder/glog v1.0.1-0.20220322161911-7365fe7f2cd1 h1:UqBrPWSYvRI2s5RtOu github.com/coder/glog v1.0.1-0.20220322161911-7365fe7f2cd1/go.mod h1:EWib/APOK0SL3dFbYqvxE3UYd8E6s1ouQ7iEp/0LWV4= github.com/coder/retry v1.3.0 h1:5lAAwt/2Cm6lVmnfBY7sOMXcBOwcwJhmV5QGSELIVWY= github.com/coder/retry v1.3.0/go.mod h1:tXuRgZgWjUnU5LZPT4lJh4ew2elUhexhlnXzrJWdyFY= -github.com/coder/tailscale v1.1.1-0.20220902164407-ae46caa65076 h1:PITEtBolloXfTMGSkL1hQSPBMT4+YJFUgjRQl5osB5k= -github.com/coder/tailscale v1.1.1-0.20220902164407-ae46caa65076/go.mod h1:MO+tWkQp2YIF3KBnnej/mQvgYccRS5Xk/IrEpZ4Z3BU= +github.com/coder/tailscale v1.1.1-0.20220905194158-291661887d25 h1:XOloZLgDkAmVBVYXSQBLY+a/Vd2c+dWRBMKNJMWSAWo= +github.com/coder/tailscale v1.1.1-0.20220905194158-291661887d25/go.mod h1:MO+tWkQp2YIF3KBnnej/mQvgYccRS5Xk/IrEpZ4Z3BU= github.com/coder/wireguard-go/tun/netstack v0.0.0-20220823170024-a78136eb0cab h1:9yEvRWXXfyKzXu8AqywCi+tFZAoqCy4wVcsXwuvZNMc= github.com/coder/wireguard-go/tun/netstack v0.0.0-20220823170024-a78136eb0cab/go.mod h1:TCJ66NtXh3urJotTdoYQOHHkyE899vOQl5TuF+WLSes= github.com/containerd/aufs v0.0.0-20200908144142-dab0cbea06f4/go.mod h1:nukgQABAEopAHvB6j7cnP5zJ+/3aVcE7hCYqvIwAHyE=