stop stats reporting at agent level

f0ssel · f0ssel · commit f90bed23f136 · 2024-06-17T19:01:36.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -86,6 +86,8 @@ type Options struct {
 	PrometheusRegistry           *prometheus.Registry
 	ReportMetadataInterval       time.Duration
 	ServiceBannerRefreshInterval time.Duration
+	ExperimentRefreshInterval    time.Duration
+	FetchExperiments             func(ctx context.Context) (codersdk.Experiments, error)
 	Syscaller                    agentproc.Syscaller
 	// ModifiedProcesses is used for testing process priority management.
 	ModifiedProcesses chan []*agentproc.Process
@@ -134,6 +136,14 @@ func New(options Options) Agent {
 			return "", nil
 		}
 	}
+	if options.FetchExperiments == nil {
+		options.FetchExperiments = func(ctx context.Context) (codersdk.Experiments, error) {
+			return codersdk.Experiments{}, nil
+		}
+	}
+	if options.ExperimentRefreshInterval == 0 {
+		options.ExperimentRefreshInterval = 5 * time.Minute
+	}
 	if options.ReportMetadataInterval == 0 {
 		options.ReportMetadataInterval = time.Second
 	}
@@ -167,6 +177,7 @@ func New(options Options) Agent {
 		environmentVariables:               options.EnvironmentVariables,
 		client:                             options.Client,
 		exchangeToken:                      options.ExchangeToken,
+		fetchExperiments:                   options.FetchExperiments,
 		filesystem:                         options.Filesystem,
 		logDir:                             options.LogDir,
 		tempDir:                            options.TempDir,
@@ -249,6 +260,10 @@ type agent struct {
 	lifecycleStates            []agentsdk.PostLifecycleRequest
 	lifecycleLastReportedIndex int // Keeps track of the last lifecycle state we successfully reported.
 
+	fetchExperiments         func(ctx context.Context) (codersdk.Experiments, error)
+	fetchExperimentsInterval time.Duration
+	experiments              atomic.Pointer[codersdk.Experiments]
+
 	network       *tailnet.Conn
 	addresses     []netip.Prefix
 	statsReporter *statsReporter
@@ -737,6 +752,28 @@ func (a *agent) fetchServiceBannerLoop(ctx context.Context, conn drpc.Conn) erro
 	}
 }
 
+// fetchExperimentsLoop fetches experiments on an interval.
+func (a *agent) fetchExperimentsLoop(ctx context.Context) error {
+	ticker := time.NewTicker(a.fetchExperimentsInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ticker.C:
+			experiments, err := a.fetchExperiments(ctx)
+			if err != nil {
+				if ctx.Err() != nil {
+					return ctx.Err()
+				}
+				a.logger.Error(ctx, "failed to update experiments", slog.Error(err))
+				return err
+			}
+			a.experiments.Store(&experiments)
+		}
+	}
+}
+
 func (a *agent) run() (retErr error) {
 	// This allows the agent to refresh it's token if necessary.
 	// For instance identity this is required, since the instance
@@ -747,6 +784,12 @@ func (a *agent) run() (retErr error) {
 	}
 	a.sessionToken.Store(&sessionToken)
 
+	exp, err := a.fetchExperiments(a.hardCtx)
+	if err != nil {
+		return xerrors.Errorf("fetch experiments: %w", err)
+	}
+	a.experiments.Store(&exp)
+
 	// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs
 	conn, err := a.client.ConnectRPC(a.hardCtx)
 	if err != nil {
@@ -856,6 +899,10 @@ func (a *agent) run() (retErr error) {
 
 	connMan.start("fetch service banner loop", gracefulShutdownBehaviorStop, a.fetchServiceBannerLoop)
 
+	connMan.start("fetch experiments loop", gracefulShutdownBehaviorStop, func(ctx context.Context, _ drpc.Conn) error {
+		return a.fetchExperimentsLoop(ctx)
+	})
+
 	connMan.start("stats report loop", gracefulShutdownBehaviorStop, func(ctx context.Context, conn drpc.Conn) error {
 		if err := networkOK.wait(ctx); err != nil {
 			return xerrors.Errorf("no network: %w", err)
diff --git a/agent/stats.go b/agent/stats.go
@@ -5,11 +5,13 @@ import (
 	"sync"
 	"time"
 
+	"go.uber.org/atomic"
 	"golang.org/x/xerrors"
 	"tailscale.com/types/netlogtype"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/agent/proto"
+	"github.com/coder/coder/v2/codersdk"
 )
 
 const maxConns = 2048
@@ -36,9 +38,10 @@ type statsReporter struct {
 	unreported   bool
 	lastInterval time.Duration
 
-	source    networkStatsSource
-	collector statsCollector
-	logger    slog.Logger
+	source      networkStatsSource
+	collector   statsCollector
+	logger      slog.Logger
+	experiments atomic.Pointer[codersdk.Experiments]
 }
 
 func newStatsReporter(logger slog.Logger, source networkStatsSource, collector statsCollector) *statsReporter {
@@ -112,6 +115,17 @@ func (s *statsReporter) reportLocked(
 	s.L.Unlock()
 	defer s.L.Lock()
 	stats := s.collector.Collect(ctx, networkStats)
+
+	// if the experiment is enabled we zero out certain session stats
+	// as we migrate to the client reporting these stats instead.
+	if s.experiments.Load().Enabled(codersdk.ExperimentWorkspaceUsage) {
+		stats.SessionCountSsh = 0
+		// TODO: More session types will be enabled as we migrate over.
+		// stats.SessionCountVscode = 0
+		// stats.SessionCountJetbrains = 0
+		// stats.SessionCountReconnectingPty = 0
+	}
+
 	resp, err := dest.UpdateStats(ctx, &proto.UpdateStatsRequest{Stats: stats})
 	if err != nil {
 		return err
diff --git a/cli/agent.go b/cli/agent.go
@@ -306,6 +306,9 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 					client.SetSessionToken(resp.SessionToken)
 					return resp.SessionToken, nil
 				},
+				FetchExperiments: func(ctx context.Context) (codersdk.Experiments, error) {
+					return client.SDK.Experiments(ctx)
+				},
 				EnvironmentVariables: environmentVariables,
 				IgnorePorts:          ignorePorts,
 				SSHMaxTimeout:        sshMaxTimeout,
