mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user