feat: send connection logs from agentfake agents (#26083)

Signed-off-by: Callum Styan <callumstyan@gmail.com>

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",

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 {

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())
+			}
+		})
+	}
+}

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