Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit afd343e

Browse files
coadlercoadler
committed
chore: update tailscale
1 parent b81d846 commit afd343e

File tree

9 files changed

+195
-108
lines changed

9 files changed

+195
-108
lines changed

‎agent/agent.go

Lines changed: 17 additions & 6 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
PostVersion(ctx context.Context, version string) 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.
@@ -338,11 +340,20 @@ func (a *agent) run(ctx context.Context) error {
338340
return xerrors.New("agent is closed")
339341
}
340342

343+
setStatInterval := func(d time.Duration) {
344+
network.SetConnStatsCallback(d, 2048,
345+
func(_, _ time.Time, virtual, _ map[netlogtype.Connection]netlogtype.Counts) {
346+
select {
347+
case a.connStatsChan <- convertAgentStats(virtual):
348+
default:
349+
a.logger.Warn(ctx, "network stat dropped")
350+
}
351+
},
352+
)
353+
}
354+
341355
// Report statistics from the created network.
342-
cl, err := a.client.ReportStats(ctx, a.logger, func() *agentsdk.Stats {
343-
stats := network.ExtractTrafficStats()
344-
return convertAgentStats(stats)
345-
})
356+
cl, err := a.client.ReportStats(ctx, a.logger, a.connStatsChan, setStatInterval)
346357
if err != nil {
347358
a.logger.Error(ctx, "report stats", slog.Error(err))
348359
} else {

‎agent/agent_test.go

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1200,28 +1200,27 @@ func (c *client) Listen(_ context.Context) (net.Conn, error) {
12001200
return clientConn, nil
12011201
}
12021202

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

12071207
go func() {
12081208
defer close(doneCh)
12091209

1210-
t := time.NewTicker(500 * time.Millisecond)
1211-
defer t.Stop()
1210+
setInterval(500 * time.Millisecond)
12121211
for {
12131212
select {
12141213
case <-ctx.Done():
12151214
return
1216-
case <-t.C:
1217-
}
1218-
select {
1219-
case c.statsChan <- stats():
1220-
case <-ctx.Done():
1221-
return
1222-
default:
1223-
// We don't want to send old stats.
1224-
continue
1215+
case stat := <-statsChan:
1216+
select {
1217+
case c.statsChan <- stat:
1218+
case <-ctx.Done():
1219+
return
1220+
default:
1221+
// We don't want to send old stats.
1222+
continue
1223+
}
12251224
}
12261225
}
12271226
}()

‎cli/vscodessh.go

Lines changed: 28 additions & 18 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,69 @@ func vscodeSSH() *cobra.Command {
117119
break
118120
}
119121
}
122+
120123
agentConn, err := client.DialWorkspaceAgent(ctx, agent.ID, &codersdk.DialWorkspaceAgentOptions{
121124
EnableTrafficStats: true,
122125
})
123126
if err != nil {
124127
return xerrors.Errorf("dial workspace agent: %w", err)
125128
}
126129
defer agentConn.Close()
130+
127131
agentConn.AwaitReachable(ctx)
128132
rawSSH, err := agentConn.SSH(ctx)
129133
if err != nil {
130134
return err
131135
}
132136
defer rawSSH.Close()
137+
133138
// Copy SSH traffic over stdio.
134139
go func() {
135140
_, _ = io.Copy(cmd.OutOrStdout(), rawSSH)
136141
}()
137142
go func() {
138143
_, _ = io.Copy(rawSSH, cmd.InOrStdin())
139144
}()
145+
140146
// The VS Code extension obtains the PID of the SSH process to
141147
// read the file below which contains network information to display.
142148
//
143149
// We get the parent PID because it's assumed `ssh` is calling this
144150
// command via the ProxyCommand SSH option.
145151
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:
154-
}
155-
stats, err := collectNetworkStats(ctx, agentConn, lastCollected)
152+
153+
statsErrChan := make(chan error, 1)
154+
agentConn.SetConnStatsCallback(networkInfoInterval, 2048, func(start, end time.Time, virtual, _ map[netlogtype.Connection]netlogtype.Counts) {
155+
stats, err := collectNetworkStats(ctx, agentConn, start, end, virtual)
156156
if err != nil {
157-
return err
157+
statsErrChan <- err
158+
return
158159
}
160+
159161
rawStats, err := json.Marshal(stats)
160162
if err != nil {
161-
return err
163+
statsErrChan <- err
164+
return
162165
}
163166
err = afero.WriteFile(fs, networkInfoFilePath, rawStats, 0600)
164167
if err != nil {
165-
return err
168+
statsErrChan <- err
169+
return
166170
}
167-
lastCollected = time.Now()
171+
})
172+
173+
select {
174+
case <-ctx.Done():
175+
return nil
176+
case err := <-statsErrChan:
177+
return err
168178
}
169179
},
170180
}
171181
cmd.Flags().StringVarP(&networkInfoDir, "network-info-dir", "", "", "Specifies a directory to write network information periodically.")
172182
cmd.Flags().StringVarP(&sessionTokenFile, "session-token-file", "", "", "Specifies a file that contains a session token.")
173183
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.")
184+
cmd.Flags().DurationVarP(&networkInfoInterval, "network-info-interval", "", 5*time.Second, "Specifies the interval to update network information.")
175185
return cmd
176186
}
177187

@@ -184,7 +194,7 @@ type sshNetworkStats struct {
184194
DownloadBytesSec int64 `json:"download_bytes_sec"`
185195
}
186196

187-
func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgentConn, lastCollected time.Time) (*sshNetworkStats, error) {
197+
func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgentConn, start, end time.Time, counts map[netlogtype.Connection]netlogtype.Counts) (*sshNetworkStats, error) {
188198
latency, p2p, err := agentConn.Ping(ctx)
189199
if err != nil {
190200
return nil, err
@@ -216,13 +226,13 @@ func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgent
216226

217227
totalRx := uint64(0)
218228
totalTx := uint64(0)
219-
for _, stat := range agentConn.ExtractTrafficStats() {
229+
for _, stat := range counts {
220230
totalRx += stat.RxBytes
221231
totalTx += stat.TxBytes
222232
}
223233
// Tracking the time since last request is required because
224234
// ExtractTrafficStats() resets its counters after each call.
225-
dur := time.Since(lastCollected)
235+
dur := end.Sub(start)
226236
uploadSecs := float64(totalTx) / dur.Seconds()
227237
downloadSecs := float64(totalRx) / dur.Seconds()
228238

‎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: 30 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -368,39 +368,46 @@ 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)
377374

378-
go func() {
379-
// Immediately trigger a stats push to get the correct interval.
380-
timer := time.NewTimer(time.Nanosecond)
381-
defer timer.Stop()
375+
postStat := func(stat *Stats) {
376+
var nextInterval time.Duration
377+
for r := retry.New(100*time.Millisecond, time.Minute); r.Wait(ctx); {
378+
resp, err := c.PostStats(ctx, stat)
379+
if err != nil {
380+
if !xerrors.Is(err, context.Canceled) {
381+
log.Error(ctx, "report stats", slog.Error(err))
382+
}
383+
continue
384+
}
385+
386+
nextInterval = resp.ReportInterval
387+
break
388+
}
389+
390+
if interval != nextInterval {
391+
setInterval(nextInterval)
392+
}
393+
interval = nextInterval
394+
}
395+
396+
// Send an empty stat to get the interval.
397+
postStat(&Stats{ConnsByProto: map[string]int64{}})
382398

399+
go func() {
383400
for {
384401
select {
385402
case <-ctx.Done():
386403
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
404+
case stat, ok := <-statsChan:
405+
if !ok {
406+
return
398407
}
399408

400-
nextInterval = resp.ReportInterval
401-
break
409+
postStat(stat)
402410
}
403-
timer.Reset(nextInterval)
404411
}
405412
}()
406413

‎codersdk/workspaceagents_test.go

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"testing"
1111
"time"
1212

13+
"github.com/stretchr/testify/assert"
1314
"github.com/stretchr/testify/require"
1415
"tailscale.com/tailcfg"
1516

@@ -55,26 +56,40 @@ func TestWorkspaceAgentMetadata(t *testing.T) {
5556
func TestAgentReportStats(t *testing.T) {
5657
t.Parallel()
5758

58-
var numReports atomic.Int64
59+
var (
60+
numReports atomic.Int64
61+
numIntervalCalls atomic.Int64
62+
wantInterval = 5 * time.Millisecond
63+
)
5964
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
6065
numReports.Add(1)
6166
httpapi.Write(context.Background(), w, http.StatusOK, agentsdk.StatsResponse{
62-
ReportInterval: 5 * time.Millisecond,
67+
ReportInterval: wantInterval,
6368
})
6469
}))
6570
parsed, err := url.Parse(srv.URL)
6671
require.NoError(t, err)
6772
client := agentsdk.New(parsed)
6873

74+
assertStatInterval := func(interval time.Duration) {
75+
numIntervalCalls.Add(1)
76+
assert.Equal(t, wantInterval, interval)
77+
}
78+
79+
chanLen := 3
80+
statCh := make(chan *agentsdk.Stats, chanLen)
81+
for i := 0; i < chanLen; i++ {
82+
statCh <- &agentsdk.Stats{ConnsByProto: map[string]int64{}}
83+
}
84+
6985
ctx := context.Background()
70-
closeStream, err := client.ReportStats(ctx, slogtest.Make(t, nil), func() *agentsdk.Stats {
71-
return &agentsdk.Stats{}
72-
})
86+
closeStream, err := client.ReportStats(ctx, slogtest.Make(t, nil), statCh, assertStatInterval)
7387
require.NoError(t, err)
7488
defer closeStream.Close()
7589

7690
require.Eventually(t,
7791
func() bool { return numReports.Load() >= 3 },
7892
testutil.WaitMedium, testutil.IntervalFast,
7993
)
94+
require.Equal(t, int64(1), numIntervalCalls.Load())
8095
}

0 commit comments

Comments
 (0)