coder · mtojek · May 25, 2023 · May 17, 2023 · May 17, 2023 · May 17, 2023
diff --git a/agent/agent.go b/agent/agent.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/armon/circbuf"
 	"github.com/google/uuid"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/spf13/afero"
 	"go.uber.org/atomic"
 	"golang.org/x/exp/slices"
@@ -63,6 +64,8 @@ type Options struct {
 	SSHMaxTimeout          time.Duration
 	TailnetListenPort      uint16
 	Subsystem              codersdk.AgentSubsystem
+
+	PrometheusRegistry *prometheus.Registry
 }
 
 type Client interface {
@@ -102,6 +105,12 @@ func New(options Options) Agent {
 			return "", nil
 		}
 	}
+
+	prometheusRegistry := options.PrometheusRegistry
+	if prometheusRegistry == nil {
+		prometheusRegistry = prometheus.NewRegistry()
+	}
+
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	a := &agent{
 		tailnetListenPort:      options.TailnetListenPort,
@@ -121,6 +130,9 @@ func New(options Options) Agent {
 		connStatsChan:          make(chan *agentsdk.Stats, 1),
 		sshMaxTimeout:          options.SSHMaxTimeout,
 		subsystem:              options.Subsystem,
+
+		prometheusRegistry: prometheusRegistry,
+		metrics:            newAgentMetrics(prometheusRegistry),
 	}
 	a.init(ctx)
 	return a
@@ -165,10 +177,13 @@ type agent struct {
 	latestStat    atomic.Pointer[agentsdk.Stats]
 
 	connCountReconnectingPTY atomic.Int64
+
+	prometheusRegistry *prometheus.Registry
+	metrics            *agentMetrics
 }
 
 func (a *agent) init(ctx context.Context) {
-	sshSrv, err := agentssh.NewServer(ctx, a.logger.Named("ssh-server"), a.filesystem, a.sshMaxTimeout, "")
+	sshSrv, err := agentssh.NewServer(ctx, a.logger.Named("ssh-server"), a.prometheusRegistry, a.filesystem, a.sshMaxTimeout, "")
 	if err != nil {
 		panic(err)
 	}
@@ -983,6 +998,7 @@ func (a *agent) trackScriptLogs(ctx context.Context, reader io.Reader) (chan str
 
 func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, msg codersdk.WorkspaceAgentReconnectingPTYInit, conn net.Conn) (retErr error) {
 	defer conn.Close()
+	a.metrics.connectionsTotal.Add(1)
 
 	a.connCountReconnectingPTY.Add(1)
 	defer a.connCountReconnectingPTY.Add(-1)
@@ -1022,6 +1038,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 		// Empty command will default to the users shell!
 		cmd, err := a.sshServer.CreateCommand(ctx, msg.Command, nil)
 		if err != nil {
+			a.metrics.reconnectingPTYErrors.WithLabelValues("create_command").Add(1)
 			return xerrors.Errorf("create command: %w", err)
 		}
 		cmd.Env = append(cmd.Env, "TERM=xterm-256color")
@@ -1034,6 +1051,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 
 		ptty, process, err := pty.Start(cmd)
 		if err != nil {
+			a.metrics.reconnectingPTYErrors.WithLabelValues("start_command").Add(1)
 			return xerrors.Errorf("start command: %w", err)
 		}
 
@@ -1060,7 +1078,12 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 				if err != nil {
 					// When the PTY is closed, this is triggered.
 					// Error is typically a benign EOF, so only log for debugging.
-					logger.Debug(ctx, "unable to read pty output, command exited?", slog.Error(err))
+					if errors.Is(err, io.EOF) {
+						logger.Debug(ctx, "unable to read pty output, command exited?", slog.Error(err))
+					} else {
+						logger.Warn(ctx, "unable to read pty output, command exited?", slog.Error(err))
+						a.metrics.reconnectingPTYErrors.WithLabelValues("output_reader").Add(1)
+					}
 					break
 				}
 				part := buffer[:read]
@@ -1075,11 +1098,12 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 				for cid, conn := range rpty.activeConns {
 					_, err = conn.Write(part)
 					if err != nil {
-						logger.Debug(ctx,
+						logger.Warn(ctx,
 							"error writing to active conn",
 							slog.F("other_conn_id", cid),
 							slog.Error(err),
 						)
+						a.metrics.reconnectingPTYErrors.WithLabelValues("write").Add(1)
 					}
 				}
 				rpty.activeConnsMutex.Unlock()
@@ -1099,6 +1123,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 	if err != nil {
 		// We can continue after this, it's not fatal!
 		logger.Error(ctx, "resize", slog.Error(err))
+		a.metrics.reconnectingPTYErrors.WithLabelValues("resize").Add(1)
 	}
 	// Write any previously stored data for the TTY.
 	rpty.circularBufferMutex.RLock()
@@ -1111,6 +1136,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 	// while also holding circularBufferMutex seems dangerous.
 	_, err = conn.Write(prevBuf)
 	if err != nil {
+		a.metrics.reconnectingPTYErrors.WithLabelValues("write").Add(1)
 		return xerrors.Errorf("write buffer to conn: %w", err)
 	}
 	// Multiple connections to the same TTY are permitted.
@@ -1161,6 +1187,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 		_, err = rpty.ptty.InputWriter().Write([]byte(req.Data))
 		if err != nil {
 			logger.Warn(ctx, "write to pty", slog.Error(err))
+			a.metrics.reconnectingPTYErrors.WithLabelValues("input_writer").Add(1)
 			return nil
 		}
 		// Check if a resize needs to happen!
@@ -1171,6 +1198,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 		if err != nil {
 			// We can continue after this, it's not fatal!
 			logger.Error(ctx, "resize", slog.Error(err))
+			a.metrics.reconnectingPTYErrors.WithLabelValues("resize").Add(1)
 		}
 	}
 }
@@ -1203,7 +1231,7 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
 		var mu sync.Mutex
 		status := a.network.Status()
 		durations := []float64{}
-		ctx, cancelFunc := context.WithTimeout(ctx, 5*time.Second)
+		pingCtx, cancelFunc := context.WithTimeout(ctx, 5*time.Second)
 		defer cancelFunc()
 		for nodeID, peer := range status.Peer {
 			if !peer.Active {
@@ -1219,7 +1247,7 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
-				duration, _, _, err := a.network.Ping(ctx, addresses[0].Addr())
+				duration, _, _, err := a.network.Ping(pingCtx, addresses[0].Addr())
 				if err != nil {
 					return
 				}
@@ -1244,7 +1272,10 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
 		// Collect agent metrics.
 		// Agent metrics are changing all the time, so there is no need to perform
 		// reflect.DeepEqual to see if stats should be transferred.
-		stats.Metrics = collectMetrics()
+
+		metricsCtx, cancelFunc := context.WithTimeout(ctx, 5*time.Second)
+		defer cancelFunc()
+		stats.Metrics = a.collectMetrics(metricsCtx)
 
 		a.latestStat.Store(stats)
 

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -27,6 +27,8 @@ import (
 	"github.com/google/uuid"
 	"github.com/pion/udp"
 	"github.com/pkg/sftp"
+	"github.com/prometheus/client_golang/prometheus"
+	promgo "github.com/prometheus/client_model/go"
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -1724,7 +1726,7 @@ func (c closeFunc) Close() error {
 	return c()
 }
 
-func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Duration) (
+func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Duration, opts ...func(agent.Options) agent.Options) (
 	*codersdk.WorkspaceAgentConn,
 	*client,
 	<-chan *agentsdk.Stats,
@@ -1749,12 +1751,19 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
 		statsChan:   statsCh,
 		coordinator: coordinator,
 	}
-	closer := agent.New(agent.Options{
+
+	options := agent.Options{
 		Client:                 c,
 		Filesystem:             fs,
 		Logger:                 logger.Named("agent"),
 		ReconnectingPTYTimeout: ptyTimeout,
-	})
+	}
+
+	for _, opt := range opts {
+		options = opt(options)
+	}
+
+	closer := agent.New(options)
 	t.Cleanup(func() {
 		_ = closer.Close()
 	})
@@ -1979,3 +1988,107 @@ func tempDirUnixSocket(t *testing.T) string {
 
 	return t.TempDir()
 }
+
+func TestAgent_Metrics_SSH(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+	defer cancel()
+
+	registry := prometheus.NewRegistry()
+
+	//nolint:dogsled
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(o agent.Options) agent.Options {
+		o.PrometheusRegistry = registry
+		return o
+	})
+
+	sshClient, err := conn.SSHClient(ctx)
+	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)
+
+	expected := []agentsdk.AgentMetric{
+		{
+			Name:  "agent_reconnecting_pty_connections_total",
+			Type:  agentsdk.AgentMetricTypeCounter,
+			Value: 0,
+		},
+		{
+			Name:  "agent_sessions_total",
+			Type:  agentsdk.AgentMetricTypeCounter,
+			Value: 1,
+			Labels: []agentsdk.AgentMetricLabel{
+				{
+					Name:  "magic_type",
+					Value: "ssh",
+				},
+				{
+					Name:  "pty",
+					Value: "no",
+				},
+			},
+		},
+		{
+			Name:  "agent_ssh_server_failed_connections_total",
+			Type:  agentsdk.AgentMetricTypeCounter,
+			Value: 0,
+		},
+		{
+			Name:  "agent_ssh_server_sftp_connections_total",
+			Type:  agentsdk.AgentMetricTypeCounter,
+			Value: 0,
+		},
+		{
+			Name:  "agent_ssh_server_sftp_server_errors_total",
+			Type:  agentsdk.AgentMetricTypeCounter,
+			Value: 0,
+		},
+	}
+
+	var actual []*promgo.MetricFamily
+	assert.Eventually(t, func() bool {
+		actual, err = registry.Gather()
+		if err != nil {
+			return false
+		}
+
+		if len(expected) != len(actual) {
+			return false
+		}
+
+		return verifyCollectedMetrics(t, expected, actual)
+	}, testutil.WaitLong, testutil.IntervalFast)
+
+	require.Len(t, actual, len(expected))
+	verifyCollectedMetrics(t, expected, actual)
+
+	_ = stdin.Close()
+	err = session.Wait()
+	require.NoError(t, err)
+}
+
+func verifyCollectedMetrics(t *testing.T, expected []agentsdk.AgentMetric, actual []*promgo.MetricFamily) bool {
+	for i, e := range expected {
+		require.Equal(t, e.Name, actual[i].GetName())
+		require.Equal(t, string(e.Type), strings.ToLower(actual[i].GetType().String()))
+
+		for _, m := range actual[i].GetMetric() {
+			require.Equal(t, e.Value, m.Counter.GetValue())
+
+			if len(m.GetLabel()) > 0 {
+				for j, lbl := range m.GetLabel() {
+					require.Equal(t, e.Labels[j].Name, lbl.GetName())
+					require.Equal(t, e.Labels[j].Value, lbl.GetValue())
+				}
+			}
+			m.GetLabel()
+		}
+	}
+	return true
+}