From 89f200872b4f22fab96add50eb449a739ebee904 Mon Sep 17 00:00:00 2001 From: Callum Styan Date: Mon, 15 Jun 2026 13:35:28 -0700 Subject: [PATCH] feat: send connection logs from `agentfake` agents (#26083) Signed-off-by: Callum Styan --- enterprise/cli/exp_scaletest_agentfake.go | 29 +++- enterprise/scaletest/agentfake/agent.go | 78 ++++++++++ enterprise/scaletest/agentfake/agent_test.go | 152 +++++++++++++++++++ enterprise/scaletest/agentfake/manager.go | 7 +- 4 files changed, 260 insertions(+), 6 deletions(-) diff --git a/enterprise/cli/exp_scaletest_agentfake.go b/enterprise/cli/exp_scaletest_agentfake.go index b3ccd51629..d72058ad16 100644 --- a/enterprise/cli/exp_scaletest_agentfake.go +++ b/enterprise/cli/exp_scaletest_agentfake.go @@ -4,6 +4,7 @@ package cli import ( "os/signal" + "time" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promhttp" @@ -40,6 +41,8 @@ func (r *RootCmd) scaletestAgentFake() *serpent.Command { expectedAgentsTolerance int64 postgresURL string postgresAuth string + connReportInterval time.Duration + connReportDuration time.Duration ) cmd := &serpent.Command{ @@ -117,11 +120,13 @@ func (r *RootCmd) scaletestAgentFake() *serpent.Command { metrics := agentfake.NewMetrics(prometheus.DefaultRegisterer) mgr := agentfake.NewManager(logger, client.URL, client, db, agentfake.ManagerOptions{ - Template: template, - Owner: owner, - Metrics: metrics, - ExpectedAgents: expectedAgents, - ExpectedAgentsTolerance: expectedAgentsTolerance, + Template: template, + Owner: owner, + Metrics: metrics, + ExpectedAgents: expectedAgents, + ExpectedAgentsTolerance: expectedAgentsTolerance, + ConnectionReportInterval: connReportInterval, + ConnectionReportDuration: connReportDuration, }) defer mgr.Close() @@ -166,6 +171,20 @@ func (r *RootCmd) scaletestAgentFake() *serpent.Command { Description: "Acceptable variance around --expected-agents. Ignored when --expected-agents is 0.", Value: serpent.Int64Of(&expectedAgentsTolerance), }, + { + Flag: "connection-report-interval", + Env: "CODER_SCALETEST_AGENTFAKE_CONNECTION_REPORT_INTERVAL", + Description: "Idle gap between synthetic SSH connect events per fake agent. Zero disables connection reporting.", + Default: "30s", + Value: serpent.DurationOf(&connReportInterval), + }, + { + Flag: "connection-report-duration", + Env: "CODER_SCALETEST_AGENTFAKE_CONNECTION_REPORT_DURATION", + Description: "Synthetic SSH session length per fake agent. Zero disables connection reporting.", + Default: "5s", + Value: serpent.DurationOf(&connReportDuration), + }, { Flag: "postgres-url", Env: "CODER_PG_CONNECTION_URL", diff --git a/enterprise/scaletest/agentfake/agent.go b/enterprise/scaletest/agentfake/agent.go index 4242e81978..a2539f1621 100644 --- a/enterprise/scaletest/agentfake/agent.go +++ b/enterprise/scaletest/agentfake/agent.go @@ -62,6 +62,11 @@ type Agent struct { firstConnect chan<- time.Duration firstConnected atomic.Bool + // A zero connReportInterval or connReportDuration disables synthetic SSH + // connection reporting. + connReportInterval time.Duration + connReportDuration time.Duration + start time.Time cancel context.CancelFunc @@ -108,6 +113,15 @@ func WithFirstConnect(ch chan<- time.Duration) Option { } } +// WithConnectionReports enables periodic synthetic SSH connection reporting. +// A zero interval or duration disables reporting. +func WithConnectionReports(interval, duration time.Duration) Option { + return func(a *Agent) { + a.connReportInterval = interval + a.connReportDuration = duration + } +} + func NewAgent(logger slog.Logger, coderURL *url.URL, token string, opts ...Option) *Agent { a := &Agent{ coderURL: coderURL, @@ -225,6 +239,9 @@ func (a *Agent) connectAndServe(ctx context.Context, client rpcDialer) error { go a.runMetadata(connCtx, rpc, workspaceID, descs) } + // Bound to connCtx so the goroutine exits on reconnect, like runMetadata. + go a.runConnectionReports(connCtx, rpc) + select { case <-ctx.Done(): return nil @@ -326,6 +343,67 @@ func (a *Agent) runMetadata(ctx context.Context, rpc proto.DRPCAgentClient29, wo }, "agentfake", "runMetadata").Wait() } +// runConnectionReports emits periodic synthetic SSH sessions (CONNECT then +// DISCONNECT) via ReportConnection. Each session reuses one connection_id so +// coderd pairs the two halves onto a single connection_log row. +func (a *Agent) runConnectionReports(ctx context.Context, rpc proto.DRPCAgentClient29) { + // A zero-length session is meaningless, so a zero interval or duration + // disables reporting entirely. + if a.connReportInterval <= 0 || a.connReportDuration <= 0 { + return + } + + // Tick at the smaller of the two so neither boundary is overshot. + tick := min(a.connReportInterval, a.connReportDuration) + + var ( + openID uuid.UUID + closeAt time.Time + nextOpen = a.clock.Now().Add(a.connReportInterval) + ) + _ = a.clock.TickerFunc(ctx, tick, func() error { + now := a.clock.Now() + switch { + case openID != uuid.Nil && !now.Before(closeAt): + // A failed DISCONNECT send is non-fatal for scaletesting, so we + // ignore the result and always reset the session. + a.sendConnection(ctx, rpc, openID, proto.Connection_DISCONNECT, now) + openID = uuid.Nil + nextOpen = now.Add(a.connReportInterval) + case openID == uuid.Nil && !now.Before(nextOpen): + id := uuid.New() + closeAt = now.Add(a.connReportDuration) + if a.sendConnection(ctx, rpc, id, proto.Connection_CONNECT, now) { + openID = id + } else { + // Leave openID nil so a failed CONNECT retries next interval + // instead of desyncing the connect/disconnect pairing. + nextOpen = now.Add(a.connReportInterval) + } + } + return nil + }, "agentfake", "connectionReports").Wait() +} + +func (a *Agent) sendConnection(ctx context.Context, rpc proto.DRPCAgentClient29, id uuid.UUID, action proto.Connection_Action, now time.Time) bool { + _, err := rpc.ReportConnection(ctx, &proto.ReportConnectionRequest{ + Connection: &proto.Connection{ + Id: id[:], + Action: action, + Type: proto.Connection_SSH, + Timestamp: timestamppb.New(now), + Ip: "127.0.0.1", + }, + }) + if err != nil && ctx.Err() == nil { + a.logger.Debug(ctx, "report connection failed", + slog.F("action", action.String()), + slog.Error(err)) + return false + } + return true +} + // Close stops the agent. Safe to call multiple times. func (a *Agent) Close() { if a.cancel != nil { diff --git a/enterprise/scaletest/agentfake/agent_test.go b/enterprise/scaletest/agentfake/agent_test.go index 846a6c9428..a7caef8d05 100644 --- a/enterprise/scaletest/agentfake/agent_test.go +++ b/enterprise/scaletest/agentfake/agent_test.go @@ -153,3 +153,155 @@ func TestAgent_SendsMetadata(t *testing.T) { t.Fatalf("timed out waiting for Agent.Run to return: %v", ctx.Err()) } } + +// Assert that the fake agent emits repeating CONNECT/DISCONNECT SSH sessions, +// pairing each session's halves under one connection id and using a fresh id +// per session. +func TestAgent_ReportsConnections(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + + const ( + interval = 30 * time.Second + duration = 5 * time.Second + ) + + mClock := quartz.NewMock(t) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) + + agentID := uuid.New() + manifest := agentsdk.Manifest{ + AgentID: agentID, + WorkspaceID: uuid.New(), + } + statsCh := make(chan *agentproto.Stats, 1) + coord := tailnet.NewCoordinator(logger) + t.Cleanup(func() { _ = coord.Close() }) + dialer := agenttest.NewClient(t, logger, agentID, manifest, statsCh, coord) + t.Cleanup(dialer.Close) + + a := agentfake.NewAgent(logger, nil, "", + agentfake.WithDialer(dialer), + agentfake.WithClock(mClock), + agentfake.WithConnectionReports(interval, duration), + ) + t.Cleanup(a.Close) + + // Trap registration so the goroutine is parked on the mock clock before + // we Advance, otherwise Advance could race startup and miss the first tick. + tickerTrap := mClock.Trap().TickerFunc("agentfake", "connectionReports") + defer tickerTrap.Close() + + runCtx, cancel := context.WithCancel(ctx) + t.Cleanup(cancel) + runErr := make(chan error, 1) + go func() { runErr <- a.Run(runCtx) }() + + tickerTrap.MustWait(ctx).Release(ctx) + + // Advance one tick period (5s) per step until at least `want` reports land. + advanceUntil := func(want int) { + t.Helper() + require.Eventually(t, func() bool { + mClock.Advance(duration).MustWait(ctx) + return len(dialer.GetConnectionReports()) >= want + }, testutil.WaitShort, testutil.IntervalFast, + "expected %d connection reports", want) + } + + advanceUntil(1) + reports := dialer.GetConnectionReports() + require.GreaterOrEqual(t, len(reports), 1) + require.Equal(t, agentproto.Connection_SSH, reports[0].GetConnection().GetType()) + require.Equal(t, agentproto.Connection_CONNECT, reports[0].GetConnection().GetAction()) + firstID := reports[0].GetConnection().GetId() + require.NotEqual(t, uuid.Nil[:], firstID) + + advanceUntil(2) + reports = dialer.GetConnectionReports() + require.Equal(t, agentproto.Connection_DISCONNECT, reports[1].GetConnection().GetAction()) + require.Equal(t, firstID, reports[1].GetConnection().GetId()) + + advanceUntil(3) + reports = dialer.GetConnectionReports() + require.Equal(t, agentproto.Connection_CONNECT, reports[2].GetConnection().GetAction()) + require.NotEqual(t, firstID, reports[2].GetConnection().GetId()) + + cancel() + select { + case err := <-runErr: + require.NoError(t, err, "Agent.Run returned unexpected error") + case <-ctx.Done(): + t.Fatalf("timed out waiting for Agent.Run to return: %v", ctx.Err()) + } +} + +// Assert that a zero interval or duration disables reporting entirely. +func TestAgent_ReportsConnections_Disabled(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + interval time.Duration + duration time.Duration + }{ + {"BothZero", 0, 0}, + {"ZeroInterval", 0, 5 * time.Second}, + {"ZeroDuration", 30 * time.Second, 0}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) + + agentID := uuid.New() + manifest := agentsdk.Manifest{ + AgentID: agentID, + WorkspaceID: uuid.New(), + } + statsCh := make(chan *agentproto.Stats, 1) + coord := tailnet.NewCoordinator(logger) + t.Cleanup(func() { _ = coord.Close() }) + dialer := agenttest.NewClient(t, logger, agentID, manifest, statsCh, coord) + t.Cleanup(dialer.Close) + + a := agentfake.NewAgent(logger, nil, "", + agentfake.WithDialer(dialer), + agentfake.WithConnectionReports(tc.interval, tc.duration), + ) + t.Cleanup(a.Close) + + runCtx, cancel := context.WithCancel(ctx) + t.Cleanup(cancel) + runErr := make(chan error, 1) + go func() { runErr <- a.Run(runCtx) }() + + // Wait for lifecycle=READY so the reporting goroutine has had its + // chance to start before we assert it stayed silent. + require.Eventually(t, func() bool { + for _, state := range dialer.GetLifecycleStates() { + if state == codersdk.WorkspaceAgentLifecycleReady { + return true + } + } + return false + }, testutil.WaitShort, testutil.IntervalFast, + "agent never reported Lifecycle=ready") + + // Give any (buggy) reporting a brief window to leak through. + time.Sleep(testutil.IntervalSlow) + + require.Empty(t, dialer.GetConnectionReports(), + "expected no ReportConnection calls when reporting is disabled") + + cancel() + select { + case err := <-runErr: + require.NoError(t, err, "Agent.Run returned unexpected error") + case <-ctx.Done(): + t.Fatalf("timed out waiting for Agent.Run to return: %v", ctx.Err()) + } + }) + } +} diff --git a/enterprise/scaletest/agentfake/manager.go b/enterprise/scaletest/agentfake/manager.go index 5993d2760f..ca2d29780d 100644 --- a/enterprise/scaletest/agentfake/manager.go +++ b/enterprise/scaletest/agentfake/manager.go @@ -68,6 +68,10 @@ type ManagerOptions struct { // before enumerating. ExpectedAgents int64 ExpectedAgentsTolerance int64 + // A zero ConnectionReportInterval or ConnectionReportDuration disables + // synthetic SSH connection reporting. + ConnectionReportInterval time.Duration + ConnectionReportDuration time.Duration // Clock is used for the workspace-count polling interval. // Defaults to the real clock; override in tests with quartz.NewMock. Clock quartz.Clock @@ -149,7 +153,8 @@ func (m *Manager) Run(ctx context.Context) error { m.logger.Named("agent-"+strconv.Itoa(i)), m.coderURL, ti.Token, WithMetrics(m.opts.Metrics), - WithFirstConnect(firstConnectCh))) + WithFirstConnect(firstConnectCh), + WithConnectionReports(m.opts.ConnectionReportInterval, m.opts.ConnectionReportDuration))) } m.mu.Lock() m.agents = agents