feat: expose current agent connections by type via prometheus (#14612)

agent/agent_test.go

@@ -2531,17 +2531,17 @@ func TestAgent_Metrics_SSH(t *testing.T) {
 	err = session.Shell()
 	require.NoError(t, err)
 
-	expected := []agentsdk.AgentMetric{
+	expected := []*proto.Stats_Metric{
 		{
 			Name:  "agent_reconnecting_pty_connections_total",
-			Type:  agentsdk.AgentMetricTypeCounter,
+			Type:  proto.Stats_Metric_COUNTER,
 			Value: 0,
 		},
 		{
 			Name:  "agent_sessions_total",
-			Type:  agentsdk.AgentMetricTypeCounter,
+			Type:  proto.Stats_Metric_COUNTER,
 			Value: 1,
-			Labels: []agentsdk.AgentMetricLabel{
+			Labels: []*proto.Stats_Metric_Label{
 				{
 					Name:  "magic_type",
 					Value: "ssh",
@@ -2554,30 +2554,46 @@ func TestAgent_Metrics_SSH(t *testing.T) {
 		},
 		{
 			Name:  "agent_ssh_server_failed_connections_total",
-			Type:  agentsdk.AgentMetricTypeCounter,
+			Type:  proto.Stats_Metric_COUNTER,
 			Value: 0,
 		},
 		{
 			Name:  "agent_ssh_server_sftp_connections_total",
-			Type:  agentsdk.AgentMetricTypeCounter,
+			Type:  proto.Stats_Metric_COUNTER,
 			Value: 0,
 		},
 		{
 			Name:  "agent_ssh_server_sftp_server_errors_total",
-			Type:  agentsdk.AgentMetricTypeCounter,
+			Type:  proto.Stats_Metric_COUNTER,
 			Value: 0,
 		},
 		{
-			Name:  "coderd_agentstats_startup_script_seconds",
-			Type:  agentsdk.AgentMetricTypeGauge,
+			Name:  "coderd_agentstats_currently_reachable_peers",
+			Type:  proto.Stats_Metric_GAUGE,
 			Value: 0,
-			Labels: []agentsdk.AgentMetricLabel{
+			Labels: []*proto.Stats_Metric_Label{
 				{
-					Name:  "success",
-					Value: "true",
+					Name:  "connection_type",
+					Value: "derp",
 				},
 			},
 		},
+		{
+			Name:  "coderd_agentstats_currently_reachable_peers",
+			Type:  proto.Stats_Metric_GAUGE,
+			Value: 1,
+			Labels: []*proto.Stats_Metric_Label{
+				{
+					Name:  "connection_type",
+					Value: "p2p",
+				},
+			},
+		},
+		{
+			Name:  "coderd_agentstats_startup_script_seconds",
+			Type:  proto.Stats_Metric_GAUGE,
+			Value: 1,
+		},
 	}
 
 	var actual []*promgo.MetricFamily
@@ -2586,17 +2602,33 @@ func TestAgent_Metrics_SSH(t *testing.T) {
 		if err != nil {
 			return false
 		}
-
-		if len(expected) != len(actual) {
-			return false
+		count := 0
+		for _, m := range actual {
+			count += len(m.GetMetric())
 		}
-
-		return verifyCollectedMetrics(t, expected, actual)
+		return count == len(expected)
 	}, testutil.WaitLong, testutil.IntervalFast)
 
-	require.Len(t, actual, len(expected))
-	collected := verifyCollectedMetrics(t, expected, actual)
-	require.True(t, collected, "expected metrics were not collected")
+	i := 0
+	for _, mf := range actual {
+		for _, m := range mf.GetMetric() {
+			assert.Equal(t, expected[i].Name, mf.GetName())
+			assert.Equal(t, expected[i].Type.String(), mf.GetType().String())
+			// Value is max expected
+			if expected[i].Type == proto.Stats_Metric_GAUGE {
+				assert.GreaterOrEqualf(t, expected[i].Value, m.GetGauge().GetValue(), "expected %s to be greater than or equal to %f, got %f", expected[i].Name, expected[i].Value, m.GetGauge().GetValue())
+			} else if expected[i].Type == proto.Stats_Metric_COUNTER {
+				assert.GreaterOrEqualf(t, expected[i].Value, m.GetCounter().GetValue(), "expected %s to be greater than or equal to %f, got %f", expected[i].Name, expected[i].Value, m.GetCounter().GetValue())
+			}
+			for j, lbl := range expected[i].Labels {
+				assert.Equal(t, m.GetLabel()[j], &promgo.LabelPair{
+					Name:  &lbl.Name,
+					Value: &lbl.Value,
+				})
+			}
+			i++
+		}
+	}
 
 	_ = stdin.Close()
 	err = session.Wait()
@@ -2828,28 +2860,6 @@ func TestAgent_ManageProcessPriority(t *testing.T) {
 	})
 }
 
-func verifyCollectedMetrics(t *testing.T, expected []agentsdk.AgentMetric, actual []*promgo.MetricFamily) bool {
-	t.Helper()
-
-	for i, e := range expected {
-		assert.Equal(t, e.Name, actual[i].GetName())
-		assert.Equal(t, string(e.Type), strings.ToLower(actual[i].GetType().String()))
-
-		for _, m := range actual[i].GetMetric() {
-			assert.Equal(t, e.Value, m.Counter.GetValue())
-
-			if len(m.GetLabel()) > 0 {
-				for j, lbl := range m.GetLabel() {
-					assert.Equal(t, e.Labels[j].Name, lbl.GetName())
-					assert.Equal(t, e.Labels[j].Value, lbl.GetValue())
-				}
-			}
-			m.GetLabel()
-		}
-	}
-	return true
-}
-
 type syncWriter struct {
 	mu sync.Mutex
 	w  io.Writer