Add 'experiment' label to heartbeat connections gauge (#90)

* Add 'experiment' label to heartbeat connections gauge * Remove metric * Add tests * Comments * Undo unnecessary changes * Make decrease conditional
m-lab · Nov 16, 2022 · 44fb82b · 44fb82b
1 parent a92b36b
commit 44fb82b
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/handler/heartbeat.go b/handler/heartbeat.go
@@ -28,7 +28,6 @@ func (c *Client) Heartbeat(rw http.ResponseWriter, req *http.Request) {
 		metrics.RequestsTotal.WithLabelValues("heartbeat", err.Error()).Inc()
 		return
 	}
-	metrics.CurrentHeartbeatConnections.Inc()
 	metrics.RequestsTotal.WithLabelValues("heartbeat", "OK").Inc()
 	go c.handleHeartbeats(ws)
 }
@@ -39,11 +38,14 @@ func (c *Client) handleHeartbeats(ws *websocket.Conn) {
 	setReadDeadline(ws)
 
 	var hostname string
+	var experiment string
 	for {
 		_, message, err := ws.ReadMessage()
 		if err != nil {
 			log.Errorf("read error: %v", err)
-			metrics.CurrentHeartbeatConnections.Dec()
+			if experiment != "" {
+				metrics.CurrentHeartbeatConnections.WithLabelValues(experiment).Dec()
+			}
 			return
 		}
 		if message != nil {
@@ -59,6 +61,8 @@ func (c *Client) handleHeartbeats(ws *websocket.Conn) {
 			case hbm.Registration != nil:
 				hostname = hbm.Registration.Hostname
 				c.RegisterInstance(*hbm.Registration)
+				experiment = hbm.Registration.Experiment
+				metrics.CurrentHeartbeatConnections.WithLabelValues(experiment).Inc()
 			case hbm.Health != nil:
 				c.UpdateHealth(hostname, *hbm.Health)
 			}

diff --git a/metrics/metrics.go b/metrics/metrics.go
@@ -37,11 +37,12 @@ var (
 	//
 	// Example usage:
 	// metrics.CurrentHeartbeatConnections.Inc()
-	CurrentHeartbeatConnections = promauto.NewGauge(
+	CurrentHeartbeatConnections = promauto.NewGaugeVec(
 		prometheus.GaugeOpts{
 			Name: "locate_current_heartbeat_connections",
 			Help: "Number of currently active Heartbeat connections.",
 		},
+		[]string{"experiment"},
 	)
 
 	// PrometheusHealthCollectionDuration is a histogram that tracks the latency of the

diff --git a/metrics/metrics_test.go b/metrics/metrics_test.go
@@ -9,7 +9,7 @@ import (
 func TestLintMetrics(t *testing.T) {
 	RequestsTotal.WithLabelValues("type", "status")
 	AppEngineTotal.WithLabelValues("country")
-	CurrentHeartbeatConnections.Set(0)
+	CurrentHeartbeatConnections.WithLabelValues("experiment").Set(0)
 	PrometheusHealthCollectionDuration.WithLabelValues("code")
 	PortChecksTotal.WithLabelValues("status")
 	KubernetesRequestsTotal.WithLabelValues("status")