Skip to content

Commit 1254e7a

Browse files
kylecarbskylecarbs
authored
feat: Add speedtest command for tailnet (#3874)
1 parent 38825b9 commit 1254e7a

File tree

9 files changed

+199
-8
lines changed

9 files changed

+199
-8
lines changed

‎agent/agent.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929
"go.uber.org/atomic"
3030
gossh "golang.org/x/crypto/ssh"
3131
"golang.org/x/xerrors"
32+
"tailscale.com/net/speedtest"
3233
"tailscale.com/tailcfg"
3334

3435
"cdr.dev/slog"
@@ -58,6 +59,7 @@ var (
5859
tailnetIP = netip.MustParseAddr("fd7a:115c:a1e0:49d6:b259:b7ac:b1b2:48f4")
5960
tailnetSSHPort = 1
6061
tailnetReconnectingPTYPort = 2
62+
tailnetSpeedtestPort = 3
6163
)
6264

6365
type Options struct {
@@ -256,6 +258,23 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
256258
go a.handleReconnectingPTY(ctx, msg, conn)
257259
}
258260
}()
261+
speedtestListener, err := a.network.Listen("tcp", ":"+strconv.Itoa(tailnetSpeedtestPort))
262+
if err != nil {
263+
a.logger.Critical(ctx, "listen for speedtest", slog.Error(err))
264+
return
265+
}
266+
go func() {
267+
for {
268+
conn, err := speedtestListener.Accept()
269+
if err != nil {
270+
a.logger.Debug(ctx, "speedtest listener failed", slog.Error(err))
271+
return
272+
}
273+
go func() {
274+
_ = speedtest.ServeConn(conn)
275+
}()
276+
}
277+
}()
259278
}
260279

261280
// runCoordinator listens for nodes and updates the self-node as it changes.

‎agent/agent_test.go

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"time"
2020

2121
"golang.org/x/xerrors"
22+
"tailscale.com/net/speedtest"
2223
"tailscale.com/tailcfg"
2324

2425
scp "github.com/bramvdbogaerde/go-scp"
@@ -547,6 +548,21 @@ func TestAgent(t *testing.T) {
547548
return err == nil
548549
}, testutil.WaitMedium, testutil.IntervalFast)
549550
})
551+
552+
t.Run("Speedtest", func(t *testing.T) {
553+
t.Parallel()
554+
if testing.Short() {
555+
t.Skip("The minimum duration for a speedtest is hardcoded in Tailscale to 5s!")
556+
}
557+
derpMap := tailnettest.RunDERPAndSTUN(t)
558+
conn, _ := setupAgent(t, agent.Metadata{
559+
DERPMap: derpMap,
560+
}, 0)
561+
defer conn.Close()
562+
res, err := conn.Speedtest(speedtest.Upload, speedtest.MinDuration)
563+
require.NoError(t, err)
564+
t.Logf("%.2f MBits/s", res[len(res)-1].MBitsPerSecond())
565+
})
550566
}
551567

552568
func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {

‎agent/conn.go

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"golang.org/x/crypto/ssh"
1717
"golang.org/x/xerrors"
1818
"tailscale.com/ipn/ipnstate"
19+
"tailscale.com/net/speedtest"
1920
"tailscale.com/tailcfg"
2021

2122
"github.com/coder/coder/peer"
@@ -39,6 +40,7 @@ type Conn interface {
3940
CloseWithError(err error) error
4041
ReconnectingPTY(id string, height, width uint16, command string) (net.Conn, error)
4142
SSH() (net.Conn, error)
43+
Speedtest(direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error)
4244
SSHClient() (*ssh.Client, error)
4345
DialContext(ctx context.Context, network string, addr string) (net.Conn, error)
4446
}
@@ -77,6 +79,10 @@ func (c *WebRTCConn) SSH() (net.Conn, error) {
7779
return channel.NetConn(), nil
7880
}
7981

82+
func (*WebRTCConn) Speedtest(_ speedtest.Direction, _ time.Duration) ([]speedtest.Result, error) {
83+
return nil, xerrors.New("not implemented")
84+
}
85+
8086
// SSHClient calls SSH to create a client that uses a weak cipher
8187
// for high throughput.
8288
func (c *WebRTCConn) SSHClient() (*ssh.Client, error) {
@@ -227,6 +233,18 @@ func (c *TailnetConn) SSHClient() (*ssh.Client, error) {
227233
return ssh.NewClient(sshConn, channels, requests), nil
228234
}
229235

236+
func (c *TailnetConn) Speedtest(direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error) {
237+
speedConn, err := c.DialContextTCP(context.Background(), netip.AddrPortFrom(tailnetIP, uint16(tailnetSpeedtestPort)))
238+
if err != nil {
239+
return nil, xerrors.Errorf("dial speedtest: %w", err)
240+
}
241+
results, err := speedtest.RunClientWithConn(direction, duration, speedConn)
242+
if err != nil {
243+
return nil, xerrors.Errorf("run speedtest: %w", err)
244+
}
245+
return results, err
246+
}
247+
230248
func (c *TailnetConn) DialContext(ctx context.Context, network string, addr string) (net.Conn, error) {
231249
_, rawPort, _ := net.SplitHostPort(addr)
232250
port, _ := strconv.Atoi(rawPort)

‎cli/root.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ func Core() []*cobra.Command {
7878
schedules(),
7979
show(),
8080
ssh(),
81+
speedtest(),
8182
start(),
8283
state(),
8384
stop(),

‎cli/speedtest.go

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package cli
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"time"
7+
8+
"cdr.dev/slog"
9+
"github.com/coder/coder/cli/cliflag"
10+
"github.com/coder/coder/cli/cliui"
11+
"github.com/coder/coder/codersdk"
12+
"github.com/jedib0t/go-pretty/v6/table"
13+
"github.com/spf13/cobra"
14+
"golang.org/x/xerrors"
15+
tsspeedtest "tailscale.com/net/speedtest"
16+
)
17+
18+
func speedtest() *cobra.Command {
19+
var (
20+
reverse bool
21+
timeStr string
22+
)
23+
cmd := &cobra.Command{
24+
Annotations: workspaceCommand,
25+
Use: "speedtest <workspace>",
26+
Short: "Run a speed test from your machine to the workspace.",
27+
RunE: func(cmd *cobra.Command, args []string) error {
28+
ctx, cancel := context.WithCancel(cmd.Context())
29+
defer cancel()
30+
31+
dur, err := time.ParseDuration(timeStr)
32+
if err != nil {
33+
return err
34+
}
35+
36+
client, err := CreateClient(cmd)
37+
if err != nil {
38+
return xerrors.Errorf("create codersdk client: %w", err)
39+
}
40+
41+
workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, cmd, client, codersdk.Me, args[0], false)
42+
if err != nil {
43+
return err
44+
}
45+
46+
err = cliui.Agent(ctx, cmd.ErrOrStderr(), cliui.AgentOptions{
47+
WorkspaceName: workspace.Name,
48+
Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
49+
return client.WorkspaceAgent(ctx, workspaceAgent.ID)
50+
},
51+
})
52+
if err != nil {
53+
return xerrors.Errorf("await agent: %w", err)
54+
}
55+
conn, err := client.DialWorkspaceAgentTailnet(ctx, slog.Logger{}, workspaceAgent.ID)
56+
if err != nil {
57+
return err
58+
}
59+
defer conn.Close()
60+
_, _ = conn.Ping()
61+
dir := tsspeedtest.Download
62+
if reverse {
63+
dir = tsspeedtest.Upload
64+
}
65+
cmd.Printf("Starting a %ds %s test...\n", int(dur.Seconds()), dir)
66+
results, err := conn.Speedtest(dir, dur)
67+
if err != nil {
68+
return err
69+
}
70+
tableWriter := cliui.Table()
71+
tableWriter.AppendHeader(table.Row{"Interval", "Transfer", "Bandwidth"})
72+
for _, r := range results {
73+
if r.Total {
74+
tableWriter.AppendSeparator()
75+
}
76+
tableWriter.AppendRow(table.Row{
77+
fmt.Sprintf("%.2f-%.2f sec", r.IntervalStart.Seconds(), r.IntervalEnd.Seconds()),
78+
fmt.Sprintf("%.4f MBits", r.MegaBits()),
79+
fmt.Sprintf("%.4f Mbits/sec", r.MBitsPerSecond()),
80+
})
81+
}
82+
_, err = fmt.Fprintln(cmd.OutOrStdout(), tableWriter.Render())
83+
return err
84+
},
85+
}
86+
cliflag.BoolVarP(cmd.Flags(), &reverse, "reverse", "r", "", false,
87+
"Specifies whether to run in reverse mode where the client receives and the server sends.")
88+
cliflag.StringVarP(cmd.Flags(), &timeStr, "time", "t", "", "5s",
89+
"Specifies the duration to monitor traffic.")
90+
return cmd
91+
}

‎cli/speedtest_test.go

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package cli_test
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"cdr.dev/slog/sloggers/slogtest"
8+
"github.com/coder/coder/agent"
9+
"github.com/coder/coder/cli/clitest"
10+
"github.com/coder/coder/coderd/coderdtest"
11+
"github.com/coder/coder/codersdk"
12+
"github.com/coder/coder/pty/ptytest"
13+
"github.com/coder/coder/testutil"
14+
"github.com/stretchr/testify/assert"
15+
)
16+
17+
func TestSpeedtest(t *testing.T) {
18+
t.Parallel()
19+
if testing.Short() {
20+
t.Skip("This test takes a minimum of 5ms per a hardcoded value in Tailscale!")
21+
}
22+
client, workspace, agentToken := setupWorkspaceForAgent(t)
23+
agentClient := codersdk.New(client.URL)
24+
agentClient.SessionToken = agentToken
25+
agentCloser := agent.New(agent.Options{
26+
FetchMetadata: agentClient.WorkspaceAgentMetadata,
27+
WebRTCDialer: agentClient.ListenWorkspaceAgent,
28+
CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
29+
Logger: slogtest.Make(t, nil).Named("agent"),
30+
})
31+
defer agentCloser.Close()
32+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
33+
34+
cmd, root := clitest.New(t, "speedtest", workspace.Name)
35+
clitest.SetupConfig(t, client, root)
36+
pty := ptytest.New(t)
37+
cmd.SetOut(pty.Output())
38+
39+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
40+
defer cancel()
41+
cmdDone := tGo(t, func() {
42+
err := cmd.ExecuteContext(ctx)
43+
assert.NoError(t, err)
44+
})
45+
<-cmdDone
46+
}

‎cli/ssh_test.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ import (
3131
"github.com/coder/coder/testutil"
3232
)
3333

34-
func setupWorkspaceForSSH(t *testing.T) (*codersdk.Client, codersdk.Workspace, string) {
34+
func setupWorkspaceForAgent(t *testing.T) (*codersdk.Client, codersdk.Workspace, string) {
3535
t.Helper()
3636
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
3737
user := coderdtest.CreateFirstUser(t, client)
@@ -69,7 +69,7 @@ func TestSSH(t *testing.T) {
6969
t.Run("ImmediateExit", func(t *testing.T) {
7070
t.Parallel()
7171

72-
client, workspace, agentToken := setupWorkspaceForSSH(t)
72+
client, workspace, agentToken := setupWorkspaceForAgent(t)
7373
cmd, root := clitest.New(t, "ssh", workspace.Name)
7474
clitest.SetupConfig(t, client, root)
7575
pty := ptytest.New(t)
@@ -104,7 +104,7 @@ func TestSSH(t *testing.T) {
104104
})
105105
t.Run("Stdio", func(t *testing.T) {
106106
t.Parallel()
107-
client, workspace, agentToken := setupWorkspaceForSSH(t)
107+
client, workspace, agentToken := setupWorkspaceForAgent(t)
108108
_, _ = tGoContext(t, func(ctx context.Context) {
109109
// Run this async so the SSH command has to wait for
110110
// the build and agent to connect!
@@ -175,7 +175,7 @@ func TestSSH(t *testing.T) {
175175

176176
t.Parallel()
177177

178-
client, workspace, agentToken := setupWorkspaceForSSH(t)
178+
client, workspace, agentToken := setupWorkspaceForAgent(t)
179179

180180
agentClient := codersdk.New(client.URL)
181181
agentClient.SessionToken = agentToken

‎go.mod

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ replace github.com/tcnksm/go-httpstat => github.com/kylecarbs/go-httpstat v0.0.0
4949

5050
// There are a few minor changes we make to Tailscale that we're slowly upstreaming. Compare here:
5151
// https://github.com/tailscale/tailscale/compare/main...coder:tailscale:main
52-
replace tailscale.com => github.com/coder/tailscale v1.1.1-0.20220902164407-ae46caa65076
52+
replace tailscale.com => github.com/coder/tailscale v1.1.1-0.20220905194158-291661887d25
5353

5454
require (
5555
cdr.dev/slog v1.4.2-0.20220525200111-18dce5c2cd5f
@@ -157,7 +157,7 @@ require (
157157
k8s.io/utils v0.0.0-20220210201930-3a6ce19ff2f9
158158
nhooyr.io/websocket v1.8.7
159159
storj.io/drpc v0.0.33-0.20220622181519-9206537a4db7
160-
tailscale.com v1.26.2
160+
tailscale.com v1.30.0
161161
)
162162

163163
require (

‎go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -352,8 +352,8 @@ github.com/coder/glog v1.0.1-0.20220322161911-7365fe7f2cd1 h1:UqBrPWSYvRI2s5RtOu
352352
github.com/coder/glog v1.0.1-0.20220322161911-7365fe7f2cd1/go.mod h1:EWib/APOK0SL3dFbYqvxE3UYd8E6s1ouQ7iEp/0LWV4=
353353
github.com/coder/retry v1.3.0 h1:5lAAwt/2Cm6lVmnfBY7sOMXcBOwcwJhmV5QGSELIVWY=
354354
github.com/coder/retry v1.3.0/go.mod h1:tXuRgZgWjUnU5LZPT4lJh4ew2elUhexhlnXzrJWdyFY=
355-
github.com/coder/tailscale v1.1.1-0.20220902164407-ae46caa65076 h1:PITEtBolloXfTMGSkL1hQSPBMT4+YJFUgjRQl5osB5k=
356-
github.com/coder/tailscale v1.1.1-0.20220902164407-ae46caa65076/go.mod h1:MO+tWkQp2YIF3KBnnej/mQvgYccRS5Xk/IrEpZ4Z3BU=
355+
github.com/coder/tailscale v1.1.1-0.20220905194158-291661887d25 h1:XOloZLgDkAmVBVYXSQBLY+a/Vd2c+dWRBMKNJMWSAWo=
356+
github.com/coder/tailscale v1.1.1-0.20220905194158-291661887d25/go.mod h1:MO+tWkQp2YIF3KBnnej/mQvgYccRS5Xk/IrEpZ4Z3BU=
357357
github.com/coder/wireguard-go/tun/netstack v0.0.0-20220823170024-a78136eb0cab h1:9yEvRWXXfyKzXu8AqywCi+tFZAoqCy4wVcsXwuvZNMc=
358358
github.com/coder/wireguard-go/tun/netstack v0.0.0-20220823170024-a78136eb0cab/go.mod h1:TCJ66NtXh3urJotTdoYQOHHkyE899vOQl5TuF+WLSes=
359359
github.com/containerd/aufs v0.0.0-20200908144142-dab0cbea06f4/go.mod h1:nukgQABAEopAHvB6j7cnP5zJ+/3aVcE7hCYqvIwAHyE=

0 commit comments

Comments
 (0)