chore: update tailscale

coder · coadler · Feb 10, 2023 · Feb 7, 2023 · Feb 7, 2023 · Feb 7, 2023
commit afd343e66e38b8e03063189da021afc60fa87af7
diff --git a/agent/agent.go b/agent/agent.go
@@ -72,7 +72,7 @@ type Options struct {
 type Client interface {
 	Metadata(ctx context.Context) (agentsdk.Metadata, error)
 	Listen(ctx context.Context) (net.Conn, error)
-	ReportStats(ctx context.Context, log slog.Logger, stats func() *agentsdk.Stats) (io.Closer, error)
+	ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
 	PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
 	PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error
 	PostVersion(ctx context.Context, version string) error
@@ -112,6 +112,7 @@ func New(options Options) io.Closer {
 		logDir:                 options.LogDir,
 		tempDir:                options.TempDir,
 		lifecycleUpdate:        make(chan struct{}, 1),
+		connStatsChan:          make(chan *agentsdk.Stats, 1),
 	}
 	a.init(ctx)
 	return a
@@ -143,7 +144,8 @@ type agent struct {
 	lifecycleMu     sync.Mutex // Protects following.
 	lifecycleState  codersdk.WorkspaceAgentLifecycle
 
-	network *tailnet.Conn
+	network       *tailnet.Conn
+	connStatsChan chan *agentsdk.Stats
 }
 
 // runLoop attempts to start the agent in a retry loop.
@@ -338,11 +340,20 @@ func (a *agent) run(ctx context.Context) error {
 			return xerrors.New("agent is closed")
 		}
 
+		setStatInterval := func(d time.Duration) {
+			network.SetConnStatsCallback(d, 2048,
+				func(_, _ time.Time, virtual, _ map[netlogtype.Connection]netlogtype.Counts) {
+					select {
+					case a.connStatsChan <- convertAgentStats(virtual):
+					default:
+						a.logger.Warn(ctx, "network stat dropped")
+					}
+				},
+			)
+		}
+
 		// Report statistics from the created network.
-		cl, err := a.client.ReportStats(ctx, a.logger, func() *agentsdk.Stats {
-			stats := network.ExtractTrafficStats()
-			return convertAgentStats(stats)
-		})
+		cl, err := a.client.ReportStats(ctx, a.logger, a.connStatsChan, setStatInterval)
 		if err != nil {
 			a.logger.Error(ctx, "report stats", slog.Error(err))
 		} else {

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -1200,28 +1200,27 @@ func (c *client) Listen(_ context.Context) (net.Conn, error) {
 	return clientConn, nil
 }
 
-func (c *client) ReportStats(ctx context.Context, _ slog.Logger, stats func() *agentsdk.Stats) (io.Closer, error) {
+func (c *client) ReportStats(ctx context.Context, _ slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error) {
 	doneCh := make(chan struct{})
 	ctx, cancel := context.WithCancel(ctx)
 
 	go func() {
 		defer close(doneCh)
 
-		t := time.NewTicker(500 * time.Millisecond)
-		defer t.Stop()
+		setInterval(500 * time.Millisecond)
 		for {
 			select {
 			case <-ctx.Done():
 				return
-			case <-t.C:
-			}
-			select {
-			case c.statsChan <- stats():
-			case <-ctx.Done():
-				return
-			default:
-				// We don't want to send old stats.
-				continue
+			case stat := <-statsChan:
+				select {
+				case c.statsChan <- stat:
+				case <-ctx.Done():
+					return
+				default:
+					// We don't want to send old stats.
+					continue
+				}
 			}
 		}
 	}()

diff --git a/cli/vscodessh.go b/cli/vscodessh.go
@@ -16,6 +16,7 @@ import (
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 	"tailscale.com/tailcfg"
+	"tailscale.com/types/netlogtype"
 
 	"github.com/coder/coder/codersdk"
 )
@@ -92,6 +93,7 @@ func vscodeSSH() *cobra.Command {
 			if err != nil {
 				return xerrors.Errorf("find workspace: %w", err)
 			}
+
 			var agent codersdk.WorkspaceAgent
 			var found bool
 			for _, resource := range workspace.LatestBuild.Resources {
@@ -117,61 +119,69 @@ func vscodeSSH() *cobra.Command {
 					break
 				}
 			}
+
 			agentConn, err := client.DialWorkspaceAgent(ctx, agent.ID, &codersdk.DialWorkspaceAgentOptions{
 				EnableTrafficStats: true,
 			})
 			if err != nil {
 				return xerrors.Errorf("dial workspace agent: %w", err)
 			}
 			defer agentConn.Close()
+
 			agentConn.AwaitReachable(ctx)
 			rawSSH, err := agentConn.SSH(ctx)
 			if err != nil {
 				return err
 			}
 			defer rawSSH.Close()
+
 			// Copy SSH traffic over stdio.
 			go func() {
 				_, _ = io.Copy(cmd.OutOrStdout(), rawSSH)
 			}()
 			go func() {
 				_, _ = io.Copy(rawSSH, cmd.InOrStdin())
 			}()
+
 			// The VS Code extension obtains the PID of the SSH process to
 			// read the file below which contains network information to display.
 			//
 			// We get the parent PID because it's assumed `ssh` is calling this
 			// command via the ProxyCommand SSH option.
 			networkInfoFilePath := filepath.Join(networkInfoDir, fmt.Sprintf("%d.json", os.Getppid()))
-			ticker := time.NewTicker(networkInfoInterval)
-			defer ticker.Stop()
-			lastCollected := time.Now()
-			for {
-				select {
-				case <-ctx.Done():
-					return nil
-				case <-ticker.C:
-				}
-				stats, err := collectNetworkStats(ctx, agentConn, lastCollected)
+
+			statsErrChan := make(chan error, 1)
+			agentConn.SetConnStatsCallback(networkInfoInterval, 2048, func(start, end time.Time, virtual, _ map[netlogtype.Connection]netlogtype.Counts) {
+				stats, err := collectNetworkStats(ctx, agentConn, start, end, virtual)
 				if err != nil {
-					return err
+					statsErrChan <- err
+					return
 				}
+
 				rawStats, err := json.Marshal(stats)
 				if err != nil {
-					return err
+					statsErrChan <- err
+					return
 				}
 				err = afero.WriteFile(fs, networkInfoFilePath, rawStats, 0600)
 				if err != nil {
-					return err
+					statsErrChan <- err
+					return
 				}
-				lastCollected = time.Now()
+			})
+
+			select {
+			case <-ctx.Done():
+				return nil
+			case err := <-statsErrChan:
+				return err
 			}
 		},
 	}
 	cmd.Flags().StringVarP(&networkInfoDir, "network-info-dir", "", "", "Specifies a directory to write network information periodically.")
 	cmd.Flags().StringVarP(&sessionTokenFile, "session-token-file", "", "", "Specifies a file that contains a session token.")
 	cmd.Flags().StringVarP(&urlFile, "url-file", "", "", "Specifies a file that contains the Coder URL.")
-	cmd.Flags().DurationVarP(&networkInfoInterval, "network-info-interval", "", 3*time.Second, "Specifies the interval to update network information.")
+	cmd.Flags().DurationVarP(&networkInfoInterval, "network-info-interval", "", 5*time.Second, "Specifies the interval to update network information.")
 	return cmd
 }
 
@@ -184,7 +194,7 @@ type sshNetworkStats struct {
 	DownloadBytesSec int64              `json:"download_bytes_sec"`
 }
 
-func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgentConn, lastCollected time.Time) (*sshNetworkStats, error) {
+func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgentConn, start, end time.Time, counts map[netlogtype.Connection]netlogtype.Counts) (*sshNetworkStats, error) {
 	latency, p2p, err := agentConn.Ping(ctx)
 	if err != nil {
 		return nil, err
@@ -216,13 +226,13 @@ func collectNetworkStats(ctx context.Context, agentConn *codersdk.WorkspaceAgent
 
 	totalRx := uint64(0)
 	totalTx := uint64(0)
-	for _, stat := range agentConn.ExtractTrafficStats() {
+	for _, stat := range counts {
 		totalRx += stat.RxBytes
 		totalTx += stat.TxBytes
 	}
 	// Tracking the time since last request is required because
 	// ExtractTrafficStats() resets its counters after each call.
-	dur := time.Since(lastCollected)
+	dur := end.Sub(start)
 	uploadSecs := float64(totalTx) / dur.Seconds()
 	downloadSecs := float64(totalRx) / dur.Seconds()
 

diff --git a/coderd/wsconncache/wsconncache_test.go b/coderd/wsconncache/wsconncache_test.go
@@ -214,7 +214,7 @@ func (c *client) Listen(_ context.Context) (net.Conn, error) {
 	return clientConn, nil
 }
 
-func (*client) ReportStats(_ context.Context, _ slog.Logger, _ func() *agentsdk.Stats) (io.Closer, error) {
+func (*client) ReportStats(_ context.Context, _ slog.Logger, _ <-chan *agentsdk.Stats, _ func(time.Duration)) (io.Closer, error) {
 	return io.NopCloser(strings.NewReader("")), nil
 }
 

diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -368,39 +368,46 @@ func (c *Client) AuthAzureInstanceIdentity(ctx context.Context) (AuthenticateRes
 
 // ReportStats begins a stat streaming connection with the Coder server.
 // It is resilient to network failures and intermittent coderd issues.
-func (c *Client) ReportStats(
-	ctx context.Context,
-	log slog.Logger,
-	getStats func() *Stats,
-) (io.Closer, error) {
+func (c *Client) ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *Stats, setInterval func(time.Duration)) (io.Closer, error) {
+	var interval time.Duration
 	ctx, cancel := context.WithCancel(ctx)
 
-	go func() {
-		// Immediately trigger a stats push to get the correct interval.
-		timer := time.NewTimer(time.Nanosecond)
-		defer timer.Stop()
+	postStat := func(stat *Stats) {
+		var nextInterval time.Duration
+		for r := retry.New(100*time.Millisecond, time.Minute); r.Wait(ctx); {
+			resp, err := c.PostStats(ctx, stat)
+			if err != nil {
+				if !xerrors.Is(err, context.Canceled) {
+					log.Error(ctx, "report stats", slog.Error(err))
+				}
+				continue
+			}
+
+			nextInterval = resp.ReportInterval
+			break
+		}
+
+		if interval != nextInterval {
+			setInterval(nextInterval)
+		}
+		interval = nextInterval
+	}
+
+	// Send an empty stat to get the interval.
+	postStat(&Stats{ConnsByProto: map[string]int64{}})
 
+	go func() {
 		for {
 			select {
 			case <-ctx.Done():
 				return
-			case <-timer.C:
-			}
-
-			var nextInterval time.Duration
-			for r := retry.New(100*time.Millisecond, time.Minute); r.Wait(ctx); {
-				resp, err := c.PostStats(ctx, getStats())
-				if err != nil {
-					if !xerrors.Is(err, context.Canceled) {
-						log.Error(ctx, "report stats", slog.Error(err))
-					}
-					continue
+			case stat, ok := <-statsChan:
+				if !ok {
+					return
 				}
 
-				nextInterval = resp.ReportInterval
-				break
+				postStat(stat)
 			}
-			timer.Reset(nextInterval)
 		}
 	}()
 

diff --git a/codersdk/workspaceagents_test.go b/codersdk/workspaceagents_test.go
@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"tailscale.com/tailcfg"
 
@@ -55,26 +56,40 @@ func TestWorkspaceAgentMetadata(t *testing.T) {
 func TestAgentReportStats(t *testing.T) {
 	t.Parallel()
 
-	var numReports atomic.Int64
+	var (
+		numReports       atomic.Int64
+		numIntervalCalls atomic.Int64
+		wantInterval     = 5 * time.Millisecond
+	)
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		numReports.Add(1)
 		httpapi.Write(context.Background(), w, http.StatusOK, agentsdk.StatsResponse{
-			ReportInterval: 5 * time.Millisecond,
+			ReportInterval: wantInterval,
 		})
 	}))
 	parsed, err := url.Parse(srv.URL)
 	require.NoError(t, err)
 	client := agentsdk.New(parsed)
 
+	assertStatInterval := func(interval time.Duration) {
+		numIntervalCalls.Add(1)
+		assert.Equal(t, wantInterval, interval)
+	}
+
+	chanLen := 3
+	statCh := make(chan *agentsdk.Stats, chanLen)
+	for i := 0; i < chanLen; i++ {
+		statCh <- &agentsdk.Stats{ConnsByProto: map[string]int64{}}
+	}
+
 	ctx := context.Background()
-	closeStream, err := client.ReportStats(ctx, slogtest.Make(t, nil), func() *agentsdk.Stats {
-		return &agentsdk.Stats{}
-	})
+	closeStream, err := client.ReportStats(ctx, slogtest.Make(t, nil), statCh, assertStatInterval)
 	require.NoError(t, err)
 	defer closeStream.Close()
 
 	require.Eventually(t,
 		func() bool { return numReports.Load() >= 3 },
 		testutil.WaitMedium, testutil.IntervalFast,
 	)
+	require.Equal(t, int64(1), numIntervalCalls.Load())
 }