From 20d678b886a1b5fc26c71bde05d9e6da9ef0f63c Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Thu, 4 Jun 2026 21:16:26 +0300 Subject: [PATCH] fix(agent): install connstats callback at statsReporter creation (#25819) The stats reporter only installed the connstats callback on the TUN device after the report loop negotiated an interval with the server. Traffic that flowed before that point (e.g. an SSH handshake) was silently dropped because the TUN wrapper's stats.Load() returned nil. We now install the connstats callback immediately and we no longer re-install the callback every interval unless the interval changed. Fixes flaky TestAgent_Stats_SSH, TestAgent_Stats_ReconnectingPTY, and TestAgent_Stats_Magic by ensuring the connstats callback is always installed before network traffic can flow. Closes coder/internal#505 Closes CODAGT-517 --- agent/agent.go | 11 +++- agent/agent_test.go | 119 +++++++++++++++++++++++++++-------- agent/agenttest/client.go | 39 +++++++++++- agent/stats.go | 34 +++++++--- agent/stats_internal_test.go | 6 +- 5 files changed, 167 insertions(+), 42 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index 33ad83c804..61b116de6b 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -117,6 +117,9 @@ type Options struct { ContextConfig agentcontextconfig.Config // DERPTLSConfig is an optional TLS config for DERP connections. DERPTLSConfig *tls.Config + // StatsReportInterval is the interval for the connstats callback + // installed at statsReporter creation. + StatsReportInterval time.Duration } type Client interface { @@ -183,6 +186,10 @@ func New(options Options) Agent { options.Execer = agentexec.DefaultExecer } + if options.StatsReportInterval == 0 { + options.StatsReportInterval = DefaultStatsReportInterval + } + if options.ListeningPortsGetter == nil { options.ListeningPortsGetter = &osListeningPortsGetter{ cacheDuration: 1 * time.Second, @@ -216,6 +223,7 @@ func New(options Options) Agent { ignorePorts: maps.Clone(options.IgnorePorts), }, reportMetadataInterval: options.ReportMetadataInterval, + statsReportInterval: options.StatsReportInterval, announcementBannersRefreshInterval: options.ServiceBannerRefreshInterval, sshMaxTimeout: options.SSHMaxTimeout, subsystems: options.Subsystems, @@ -289,6 +297,7 @@ type agent struct { // values. Callers that need secrets must explicitly load this. secrets atomic.Pointer[[]agentsdk.WorkspaceSecret] reportMetadataInterval time.Duration + statsReportInterval time.Duration scriptRunner *agentscripts.Runner announcementBanners atomic.Pointer[[]codersdk.BannerConfig] // announcementBanners is atomic because it is periodically updated. announcementBannersRefreshInterval time.Duration @@ -1500,7 +1509,7 @@ func (a *agent) createOrUpdateNetwork(manifestOK, networkOK *checkpoint) func(co closing := a.closing if !closing { a.network = network - a.statsReporter = newStatsReporter(a.logger, network, a) + a.statsReporter = newStatsReporter(a.logger, network, a, a.statsReportInterval) } a.closeMutex.Unlock() if closing { diff --git a/agent/agent_test.go b/agent/agent_test.go index df237e6443..dfa60e021e 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -148,33 +148,11 @@ func TestAgent_Stats_SSH(t *testing.T) { err = session.Shell() require.NoError(t, err) - var s *proto.Stats - // We are looking for four different stats to be reported. They might not all - // arrive at the same time, so we loop until we've seen them all. - var connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountSSHSeen bool - require.Eventuallyf(t, func() bool { - var ok bool - s, ok = <-stats - if !ok { - return false - } - if s.ConnectionCount > 0 { - connectionCountSeen = true - } - if s.RxBytes > 0 { - rxBytesSeen = true - } - if s.TxBytes > 0 { - txBytesSeen = true - } - if s.SessionCountSsh == 1 { - sessionCountSSHSeen = true - } - return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountSSHSeen - }, testutil.WaitLong, testutil.IntervalFast, - "never saw all stats: %+v, saw connectionCount: %t, rxBytes: %t, txBytes: %t, sessionCountSsh: %t", - s, connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountSSHSeen, - ) + // Generate SSH traffic so the connstats window sees the session. + _, err = stdin.Write([]byte("echo test\n")) + require.NoError(t, err) + + assertSSHStats(t, stats) _, err = stdin.Write([]byte("exit 0\n")) require.NoError(t, err, "writing exit to stdin") _ = stdin.Close() @@ -182,6 +160,92 @@ func TestAgent_Stats_SSH(t *testing.T) { require.NoError(t, err, "waiting for session to exit") }) } + + // Regression test for CODAGT-517: the barrier blocks reportLoop's + // initial UpdateStats, so on unfixed code the connstats callback is + // never installed and handshake traffic is lost. On fixed code the + // callback is installed at creation, so traffic is captured. + t.Run("StatsCallbackRace", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancel() + + barrier := make(chan struct{}) + + //nolint:dogsled + conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, + func(c *agenttest.Client, _ *agent.Options) { + c.SetUpdateStatsOverride(func( + ctx context.Context, + req *proto.UpdateStatsRequest, + next func(context.Context, *proto.UpdateStatsRequest) (*proto.UpdateStatsResponse, error), + ) (*proto.UpdateStatsResponse, error) { + if req.Stats == nil { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-barrier: + } + } + return next(ctx, req) + }) + }, + ) + + // Connect SSH while the barrier holds reportLoop blocked. + sshClient, err := conn.SSHClientOnPort(ctx, workspacesdk.AgentStandardSSHPort) + require.NoError(t, err) + defer sshClient.Close() + session, err := sshClient.NewSession() + require.NoError(t, err) + defer session.Close() + stdin, err := session.StdinPipe() + require.NoError(t, err) + err = session.Shell() + require.NoError(t, err) + + // Shell must be idle so the only traffic is the SSH handshake. + + close(barrier) + + assertSSHStats(t, stats) + _, err = stdin.Write([]byte("exit 0\n")) + require.NoError(t, err, "writing exit to stdin") + _ = stdin.Close() + err = session.Wait() + require.NoError(t, err, "waiting for session to exit") + }) +} + +// assertSSHStats waits for ConnectionCount, RxBytes, TxBytes, and +// SessionCountSsh to be nonzero on the stats channel. +func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) { + t.Helper() + var connectionCountSeen, rxBytesSeen, txBytesSeen, sessionCountSSHSeen bool + require.Eventuallyf(t, func() bool { + s, ok := <-stats + if !ok { + return false + } + t.Logf("got stats: ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountSsh=%d", + s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCountSsh) + if s.ConnectionCount > 0 { + connectionCountSeen = true + } + if s.RxBytes > 0 { + rxBytesSeen = true + } + if s.TxBytes > 0 { + txBytesSeen = true + } + if s.SessionCountSsh == 1 { + sessionCountSSHSeen = true + } + return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountSSHSeen + }, testutil.WaitLong, testutil.IntervalFast, + "never saw all SSH stats", + ) } func TestAgent_Stats_ReconnectingPTY(t *testing.T) { @@ -3851,6 +3915,7 @@ func setupAgentWithSecrets(t testing.TB, metadata agentsdk.Manifest, secrets []a Logger: logger.Named("agent"), ReconnectingPTYTimeout: ptyTimeout, EnvironmentVariables: map[string]string{}, + StatsReportInterval: agenttest.StatsInterval, } for _, opt := range opts { diff --git a/agent/agenttest/client.go b/agent/agenttest/client.go index 474469d7ff..24fa036119 100644 --- a/agent/agenttest/client.go +++ b/agent/agenttest/client.go @@ -32,7 +32,8 @@ import ( "github.com/coder/websocket" ) -const statsInterval = 500 * time.Millisecond +// StatsInterval is the report interval returned by FakeAgentAPI.UpdateStats. +const StatsInterval = 500 * time.Millisecond func NewClient(t testing.TB, logger slog.Logger, @@ -128,6 +129,17 @@ func (c *Client) RefreshToken(context.Context) error { return nil } +// SetUpdateStatsOverride sets a function that wraps UpdateStats calls. +// The provided function receives a next callback for the default behavior. +func (c *Client) SetUpdateStatsOverride(fn func( + ctx context.Context, + req *agentproto.UpdateStatsRequest, + next func(context.Context, *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error), +) (*agentproto.UpdateStatsResponse, error), +) { + c.fakeAgentAPI.SetUpdateStatsOverride(fn) +} + func (c *Client) GetNumRefreshTokenCalls() int { c.mu.Lock() defer c.mu.Unlock() @@ -246,6 +258,11 @@ type FakeAgentAPI struct { subAgentDisplayApps map[uuid.UUID][]agentproto.CreateSubAgentRequest_DisplayApp subAgentApps map[uuid.UUID][]*agentproto.CreateSubAgentRequest_App + updateStatsOverride func( + ctx context.Context, + req *agentproto.UpdateStatsRequest, + next func(context.Context, *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error), + ) (*agentproto.UpdateStatsResponse, error) getAnnouncementBannersFunc func() ([]codersdk.BannerConfig, error) getResourcesMonitoringConfigurationFunc func() (*agentproto.GetResourcesMonitoringConfigurationResponse, error) pushResourcesMonitoringUsageFunc func(*agentproto.PushResourcesMonitoringUsageRequest) (*agentproto.PushResourcesMonitoringUsageResponse, error) @@ -320,8 +337,26 @@ func (f *FakeAgentAPI) PushResourcesMonitoringUsage(_ context.Context, req *agen return f.pushResourcesMonitoringUsageFunc(req) } +func (f *FakeAgentAPI) SetUpdateStatsOverride(fn func( + ctx context.Context, + req *agentproto.UpdateStatsRequest, + next func(context.Context, *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error), +) (*agentproto.UpdateStatsResponse, error), +) { + f.Lock() + defer f.Unlock() + f.updateStatsOverride = fn +} + func (f *FakeAgentAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error) { f.logger.Debug(ctx, "update stats called", slog.F("req", req)) + if f.updateStatsOverride != nil { + return f.updateStatsOverride(ctx, req, f.updateStatsDefault) + } + return f.updateStatsDefault(ctx, req) +} + +func (f *FakeAgentAPI) updateStatsDefault(ctx context.Context, req *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error) { // empty request is sent to get the interval; but our tests don't want empty stats requests if req.Stats != nil { select { @@ -331,7 +366,7 @@ func (f *FakeAgentAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateSt // OK! } } - return &agentproto.UpdateStatsResponse{ReportInterval: durationpb.New(statsInterval)}, nil + return &agentproto.UpdateStatsResponse{ReportInterval: durationpb.New(StatsInterval)}, nil } func (f *FakeAgentAPI) GetLifecycleStates() []codersdk.WorkspaceAgentLifecycle { diff --git a/agent/stats.go b/agent/stats.go index 3df0fd44df..1989ff4fed 100644 --- a/agent/stats.go +++ b/agent/stats.go @@ -42,13 +42,22 @@ type statsReporter struct { logger slog.Logger } -func newStatsReporter(logger slog.Logger, source networkStatsSource, collector statsCollector) *statsReporter { - return &statsReporter{ - Cond: sync.NewCond(&sync.Mutex{}), - logger: logger, - source: source, - collector: collector, +// DefaultStatsReportInterval matches coderd.Options.AgentStatsRefreshInterval. +const DefaultStatsReportInterval = 5 * time.Minute + +func newStatsReporter(logger slog.Logger, source networkStatsSource, collector statsCollector, interval time.Duration) *statsReporter { + s := &statsReporter{ + Cond: sync.NewCond(&sync.Mutex{}), + logger: logger, + source: source, + collector: collector, + lastInterval: interval, } + // Install the callback immediately so traffic is tracked before + // reportLoop starts. reportLoop replaces it only if the + // server-negotiated interval differs. + source.SetConnStatsCallback(interval, maxConns, s.callback) + return s } func (s *statsReporter) callback(_, _ time.Time, virtual, _ map[netlogtype.Connection]netlogtype.Counts) { @@ -67,8 +76,10 @@ func (s *statsReporter) callback(_, _ time.Time, virtual, _ map[netlogtype.Conne s.Broadcast() } -// reportLoop programs the source (tailnet.Conn) to send it stats via the -// callback, then reports them to the dest. +// reportLoop reports collected stats to the server. +// +// The connstats callback is already installed by newStatsReporter; +// reportLoop only replaces it if the server returns a different interval. // // It's intended to be called within the larger retry loop that establishes a // connection to the agent API, then passes that connection to go routines like @@ -80,8 +91,11 @@ func (s *statsReporter) reportLoop(ctx context.Context, dest statsDest) error { if err != nil { return xerrors.Errorf("initial update: %w", err) } - s.lastInterval = resp.ReportInterval.AsDuration() - s.source.SetConnStatsCallback(s.lastInterval, maxConns, s.callback) + interval := resp.ReportInterval.AsDuration() + if interval != s.lastInterval { + s.lastInterval = interval + s.source.SetConnStatsCallback(s.lastInterval, maxConns, s.callback) + } // use a separate goroutine to monitor the context so that we notice immediately, rather than // waiting for the next callback (which might never come if we are closing!) diff --git a/agent/stats_internal_test.go b/agent/stats_internal_test.go index e35fa9d3e2..f0854659fc 100644 --- a/agent/stats_internal_test.go +++ b/agent/stats_internal_test.go @@ -23,7 +23,9 @@ func TestStatsReporter(t *testing.T) { fSource := newFakeNetworkStatsSource(ctx, t) fCollector := newFakeCollector(t) fDest := newFakeStatsDest() - uut := newStatsReporter(logger, fSource, fCollector) + uut := newStatsReporter(logger, fSource, fCollector, DefaultStatsReportInterval) + + _ = testutil.TryReceive(ctx, t, fSource.period) // drain construction-time install loopErr := make(chan error, 1) loopCtx, loopCancel := context.WithCancel(ctx) @@ -157,7 +159,7 @@ func newFakeNetworkStatsSource(ctx context.Context, t testing.TB) *fakeNetworkSt f := &fakeNetworkStatsSource{ ctx: ctx, t: t, - period: make(chan time.Duration), + period: make(chan time.Duration, 1), } return f }