Skip to content

Commit 4432cd0

Browse files
coadlercoadler
authored
chore: update tailscale (coder#6091)
1 parent e6da7af commit 4432cd0

11 files changed

Lines changed: 250 additions & 184 deletions

File tree

‎agent/agent.go‎

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ type Options struct {
7272
type Client interface {
7373
Metadata(ctx context.Context) (agentsdk.Metadata, error)
7474
Listen(ctx context.Context) (net.Conn, error)
75-
ReportStats(ctx context.Context, log slog.Logger, stats func() *agentsdk.Stats) (io.Closer, error)
75+
ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
7676
PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
7777
PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error
7878
PostStartup(ctx context.Context, req agentsdk.PostStartupRequest) error
@@ -112,6 +112,7 @@ func New(options Options) io.Closer {
112112
logDir: options.LogDir,
113113
tempDir: options.TempDir,
114114
lifecycleUpdate: make(chan struct{}, 1),
115+
connStatsChan: make(chan *agentsdk.Stats, 1),
115116
}
116117
a.init(ctx)
117118
return a
@@ -143,7 +144,8 @@ type agent struct {
143144
lifecycleMu sync.Mutex // Protects following.
144145
lifecycleState codersdk.WorkspaceAgentLifecycle
145146

146-
network *tailnet.Conn
147+
network *tailnet.Conn
148+
connStatsChan chan *agentsdk.Stats
147149
}
148150

149151
// runLoop attempts to start the agent in a retry loop.
@@ -351,11 +353,20 @@ func (a *agent) run(ctx context.Context) error {
351353
return xerrors.New("agent is closed")
352354
}
353355

356+
setStatInterval := func(d time.Duration) {
357+
network.SetConnStatsCallback(d, 2048,
358+
func(_, _ time.Time, virtual, _ map[netlogtype.Connection]netlogtype.Counts) {
359+
select {
360+
case a.connStatsChan <- convertAgentStats(virtual):
361+
default:
362+
a.logger.Warn(ctx, "network stat dropped")
363+
}
364+
},
365+
)
366+
}
367+
354368
// Report statistics from the created network.
355-
cl, err := a.client.ReportStats(ctx, a.logger, func() *agentsdk.Stats {
356-
stats := network.ExtractTrafficStats()
357-
return convertAgentStats(stats)
358-
})
369+
cl, err := a.client.ReportStats(ctx, a.logger, a.connStatsChan, setStatInterval)
359370
if err != nil {
360371
a.logger.Error(ctx, "report stats", slog.Error(err))
361372
} else {
@@ -399,10 +410,9 @@ func (a *agent) trackConnGoroutine(fn func()) error {
399410

400411
func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_ *tailnet.Conn, err error) {
401412
network, err := tailnet.NewConn(&tailnet.Options{
402-
Addresses: []netip.Prefix{netip.PrefixFrom(codersdk.WorkspaceAgentIP, 128)},
403-
DERPMap: derpMap,
404-
Logger: a.logger.Named("tailnet"),
405-
EnableTrafficStats: true,
413+
Addresses: []netip.Prefix{netip.PrefixFrom(codersdk.WorkspaceAgentIP, 128)},
414+
DERPMap: derpMap,
415+
Logger: a.logger.Named("tailnet"),
406416
})
407417
if err != nil {
408418
return nil, xerrors.Errorf("create tailnet: %w", err)

‎agent/agent_test.go‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,6 @@ import (
2222
"testing"
2323
"time"
2424

25-
"golang.org/x/xerrors"
26-
"tailscale.com/net/speedtest"
27-
"tailscale.com/tailcfg"
28-
2925
scp "github.com/bramvdbogaerde/go-scp"
3026
"github.com/google/uuid"
3127
"github.com/pion/udp"
@@ -37,6 +33,9 @@ import (
3733
"golang.org/x/crypto/ssh"
3834
"golang.org/x/text/encoding/unicode"
3935
"golang.org/x/text/transform"
36+
"golang.org/x/xerrors"
37+
"tailscale.com/net/speedtest"
38+
"tailscale.com/tailcfg"
4039

4140
"cdr.dev/slog"
4241
"cdr.dev/slog/sloggers/slogtest"
@@ -53,6 +52,8 @@ func TestMain(m *testing.M) {
5352
goleak.VerifyTestMain(m)
5453
}
5554

55+
// NOTE: These tests only work when your default shell is bash for some reason.
56+
5657
func TestAgent_Stats_SSH(t *testing.T) {
5758
t.Parallel()
5859
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@@ -1153,17 +1154,16 @@ func setupAgent(t *testing.T, metadata agentsdk.Metadata, ptyTimeout time.Durati
11531154
closer := agent.New(agent.Options{
11541155
Client: c,
11551156
Filesystem: fs,
1156-
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
1157+
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
11571158
ReconnectingPTYTimeout: ptyTimeout,
11581159
})
11591160
t.Cleanup(func() {
11601161
_ = closer.Close()
11611162
})
11621163
conn, err := tailnet.NewConn(&tailnet.Options{
1163-
Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
1164-
DERPMap: metadata.DERPMap,
1165-
Logger: slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug),
1166-
EnableTrafficStats: true,
1164+
Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
1165+
DERPMap: metadata.DERPMap,
1166+
Logger: slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug),
11671167
})
11681168
require.NoError(t, err)
11691169
clientConn, serverConn := net.Pipe()
@@ -1251,28 +1251,27 @@ func (c *client) Listen(_ context.Context) (net.Conn, error) {
12511251
return clientConn, nil
12521252
}
12531253

1254-
func (c *client) ReportStats(ctx context.Context, _ slog.Logger, stats func() *agentsdk.Stats) (io.Closer, error) {
1254+
func (c *client) ReportStats(ctx context.Context, _ slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error) {
12551255
doneCh := make(chan struct{})
12561256
ctx, cancel := context.WithCancel(ctx)
12571257

12581258
go func() {
12591259
defer close(doneCh)
12601260

1261-
t := time.NewTicker(500 * time.Millisecond)
1262-
defer t.Stop()
1261+
setInterval(500 * time.Millisecond)
12631262
for {
12641263
select {
12651264
case <-ctx.Done():
12661265
return
1267-
case <-t.C:
1268-
}
1269-
select {
1270-
case c.statsChan <- stats():
1271-
case <-ctx.Done():
1272-
return
1273-
default:
1274-
// We don't want to send old stats.
1275-
continue
1266+
case stat := <-statsChan:
1267+
select {
1268+
case c.statsChan <- stat:
1269+
case <-ctx.Done():
1270+
return
1271+
default:
1272+
// We don't want to send old stats.
1273+
continue
1274+
}
12761275
}
12771276
}
12781277
}()

‎cli/vscodessh.go‎

Lines changed: 39 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/spf13/cobra"
1717
"golang.org/x/xerrors"
1818
"tailscale.com/tailcfg"
19+
"tailscale.com/types/netlogtype"
1920

2021
"github.com/coder/coder/codersdk"
2122
)
@@ -92,6 +93,7 @@ func vscodeSSH() *cobra.Command {
9293
if err != nil {
9394
return xerrors.Errorf("find workspace: %w", err)
9495
}
96+
9597
var agent codersdk.WorkspaceAgent
9698
var found bool
9799
for _, resource := range workspace.LatestBuild.Resources {
@@ -117,61 +119,78 @@ func vscodeSSH() *cobra.Command {
117119
break
118120
}
119121
}
120-
agentConn, err := client.DialWorkspaceAgent(ctx, agent.ID, &codersdk.DialWorkspaceAgentOptions{
121-
EnableTrafficStats: true,
122-
})
122+
123+
agentConn, err := client.DialWorkspaceAgent(ctx, agent.ID, &codersdk.DialWorkspaceAgentOptions{})
123124
if err != nil {
124125
return xerrors.Errorf("dial workspace agent: %w", err)
125126
}
126127
defer agentConn.Close()
128+
127129
agentConn.AwaitReachable(ctx)
128130
rawSSH, err := agentConn.SSH(ctx)
129131
if err != nil {
130132
return err
131133
}
132134
defer rawSSH.Close()
135+
133136
// Copy SSH traffic over stdio.
134137
go func() {
135138
_, _ = io.Copy(cmd.OutOrStdout(), rawSSH)
136139
}()
137140
go func() {
138141
_, _ = io.Copy(rawSSH, cmd.InOrStdin())
139142
}()
143+
140144
// The VS Code extension obtains the PID of the SSH process to
141145
// read the file below which contains network information to display.
142146
//
143147
// We get the parent PID because it's assumed `ssh` is calling this
144148
// command via the ProxyCommand SSH option.
145149
networkInfoFilePath := filepath.Join(networkInfoDir, fmt.Sprintf("%d.json", os.Getppid()))
146-
ticker := time.NewTicker(networkInfoInterval)
147-
defer ticker.Stop()
148-
lastCollected := time.Now()
149-
for {
150-
select {
151-
case <-ctx.Done():
152-
return nil
153-
case <-ticker.C:
150+
151+
statsErrChan := make(chan error, 1)
152+
cb := func(start, end time.Time, virtual, _ map[netlogtype.Connection]netlogtype.Counts) {
153+
sendErr := func(err error) {
154+
select {
155+
case statsErrChan <- err:
156+
default:
157+
}
154158
}
155-
stats, err := collectNetworkStats(ctx, agentConn, lastCollected)
159+
160+
stats, err := collectNetworkStats(ctx, agentConn, start, end, virtual)
156161
if err != nil {
157-
return err
162+
sendErr(err)
163+
return
158164
}
165+
159166
rawStats, err := json.Marshal(stats)
160167
if err != nil {
161-
return err
168+
sendErr(err)
169+
return
162170
}
163171
err = afero.WriteFile(fs, networkInfoFilePath, rawStats, 0600)
164172
if err != nil {
165-
return err
173+
sendErr(err)
174+
return
166175
}
167-
lastCollected = time.Now()
176+
}
177+
178+
now := time.Now()
179+
cb(now, now.Add(time.Nanosecond), map[netlogtype.Connection]netlogtype.Counts{}, map[netlogtype.Connection]netlogtype.Counts{})
180+
agentConn.SetConnStatsCallback(networkInfoInterval, 2048, cb)
181+
182+
select {
183+
case <-ctx.Done():
184+
return nil
185+
case err := <-statsErrChan:
186+
return err
168187
}
169188
},
170189
}
171190
cmd.Flags().StringVarP(&networkInfoDir, "network-info-dir", "", "", "Specifies a directory to write network information periodically.")
172191
cmd.Flags().StringVarP(&sessionTokenFile, "session-token-file", "", "", "Specifies a file that contains a session token.")
173192
cmd.Flags().StringVarP(&urlFile, "url-file", "", "", "Specifies a file that contains the Coder URL.")
174-
cmd.Flags().DurationVarP(&networkInfoInterval, "network-info-interval", "", 3*time.Second, "Specifies the interval to update network information.")
193+
cmd.Flags().DurationVarP(&networkInfoInterval, "network-info-interval", "", 5*time.Second, "Specifies the interval to update network information.")
175194
return cmd
176195
}
177196

@@ -184,7 +203,7 @@ type sshNetworkStats struct {
184203
DownloadBytesSec int64 `json:"download_bytes_sec"`
185204
}
186205

187-
func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgentConn, lastCollected time.Time) (*sshNetworkStats, error) {
206+
func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgentConn, start, end time.Time, counts map[netlogtype.Connection]netlogtype.Counts) (*sshNetworkStats, error) {
188207
latency, p2p, err := agentConn.Ping(ctx)
189208
if err != nil {
190209
return nil, err
@@ -216,13 +235,13 @@ func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgent
216235

217236
totalRx := uint64(0)
218237
totalTx := uint64(0)
219-
for _, stat := range agentConn.ExtractTrafficStats() {
238+
for _, stat := range counts {
220239
totalRx += stat.RxBytes
221240
totalTx += stat.TxBytes
222241
}
223242
// Tracking the time since last request is required because
224243
// ExtractTrafficStats() resets its counters after each call.
225-
dur := time.Since(lastCollected)
244+
dur := end.Sub(start)
226245
uploadSecs := float64(totalTx) / dur.Seconds()
227246
downloadSecs := float64(totalRx) / dur.Seconds()
228247

‎coderd/wsconncache/wsconncache_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,7 @@ func (c *client) Listen(_ context.Context) (net.Conn, error) {
214214
return clientConn, nil
215215
}
216216

217-
func (*client) ReportStats(_ context.Context, _ slog.Logger, _ func() *agentsdk.Stats) (io.Closer, error) {
217+
func (*client) ReportStats(_ context.Context, _ slog.Logger, _ <-chan *agentsdk.Stats, _ func(time.Duration)) (io.Closer, error) {
218218
return io.NopCloser(strings.NewReader("")), nil
219219
}
220220

‎codersdk/agentsdk/agentsdk.go‎

Lines changed: 33 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -368,44 +368,55 @@ func (c *Client) AuthAzureInstanceIdentity(ctx context.Context) (AuthenticateRes
368368

369369
// ReportStats begins a stat streaming connection with the Coder server.
370370
// It is resilient to network failures and intermittent coderd issues.
371-
func (c *Client) ReportStats(
372-
ctx context.Context,
373-
log slog.Logger,
374-
getStats func() *Stats,
375-
) (io.Closer, error) {
371+
func (c *Client) ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *Stats, setInterval func(time.Duration)) (io.Closer, error) {
372+
var interval time.Duration
376373
ctx, cancel := context.WithCancel(ctx)
374+
exited := make(chan struct{})
375+
376+
postStat := func(stat *Stats) {
377+
var nextInterval time.Duration
378+
for r := retry.New(100*time.Millisecond, time.Minute); r.Wait(ctx); {
379+
resp, err := c.PostStats(ctx, stat)
380+
if err != nil {
381+
if !xerrors.Is(err, context.Canceled) {
382+
log.Error(ctx, "report stats", slog.Error(err))
383+
}
384+
continue
385+
}
386+
387+
nextInterval = resp.ReportInterval
388+
break
389+
}
390+
391+
if nextInterval != 0 && interval != nextInterval {
392+
setInterval(nextInterval)
393+
}
394+
interval = nextInterval
395+
}
396+
397+
// Send an empty stat to get the interval.
398+
postStat(&Stats{ConnsByProto: map[string]int64{}})
377399

378400
go func() {
379-
// Immediately trigger a stats push to get the correct interval.
380-
timer := time.NewTimer(time.Nanosecond)
381-
defer timer.Stop()
401+
defer close(exited)
382402

383403
for {
384404
select {
385405
case <-ctx.Done():
386406
return
387-
case <-timer.C:
388-
}
389-
390-
var nextInterval time.Duration
391-
for r := retry.New(100*time.Millisecond, time.Minute); r.Wait(ctx); {
392-
resp, err := c.PostStats(ctx, getStats())
393-
if err != nil {
394-
if !xerrors.Is(err, context.Canceled) {
395-
log.Error(ctx, "report stats", slog.Error(err))
396-
}
397-
continue
407+
case stat, ok := <-statsChan:
408+
if !ok {
409+
return
398410
}
399411

400-
nextInterval = resp.ReportInterval
401-
break
412+
postStat(stat)
402413
}
403-
timer.Reset(nextInterval)
404414
}
405415
}()
406416

407417
return closeFunc(func() error {
408418
cancel()
419+
<-exited
409420
return nil
410421
}), nil
411422
}

0 commit comments

Comments
 (0)