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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy