coder · mtojek · Jan 29, 2024 · Jan 25, 2024 · Jan 25, 2024 · Jan 25, 2024
diff --git a/coderd/prometheusmetrics/aggregator.go b/coderd/prometheusmetrics/aggregator.go
@@ -2,6 +2,9 @@ package prometheusmetrics
 
 import (
 	"context"
+	"fmt"
+	"sort"
+	"strings"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -24,20 +27,23 @@ const (
 	loggerName = "prometheusmetrics"
 
 	sizeCollectCh = 10
-	sizeUpdateCh  = 1024
+	sizeUpdateCh  = 4096
 
 	defaultMetricsCleanupInterval = 2 * time.Minute
 )
 
+var MetricLabelValueEncoder = strings.NewReplacer("\\", "\\\\", "|", "\\|", ",", "\\,", "=", "\\=")
+
 type MetricsAggregator struct {
-	queue []annotatedMetric
+	store map[metricKey]annotatedMetric
 
 	log                    slog.Logger
 	metricsCleanupInterval time.Duration
 
 	collectCh chan (chan []prometheus.Metric)
 	updateCh  chan updateRequest
 
+	storeSizeGauge   prometheus.Gauge
 	updateHistogram  prometheus.Histogram
 	cleanupHistogram prometheus.Histogram
 }
@@ -64,17 +70,37 @@ type annotatedMetric struct {
 	expiryDate time.Time
 }
 
-var _ prometheus.Collector = new(MetricsAggregator)
+type metricKey struct {
+	username      string
+	workspaceName string
+	agentName     string
+	templateName  string
 
-func (am *annotatedMetric) is(req updateRequest, m *agentproto.Stats_Metric) bool {
-	return am.username == req.username &&
-		am.workspaceName == req.workspaceName &&
-		am.agentName == req.agentName &&
-		am.templateName == req.templateName &&
-		am.Name == m.Name &&
-		agentproto.LabelsEqual(am.Labels, m.Labels)
+	metricName string
+	labelsStr  string
 }
 
+func hashKey(req *updateRequest, m *agentproto.Stats_Metric) metricKey {
+	labelPairs := make(sort.StringSlice, 0, len(m.GetLabels()))
+	for _, label := range m.GetLabels() {
+		if label.Value == "" {
+			continue
+		}
+		labelPairs = append(labelPairs, fmt.Sprintf("%s=%s", label.Name, MetricLabelValueEncoder.Replace(label.Value)))
+	}
+	labelPairs.Sort()
+	return metricKey{
+		username:      req.username,
+		workspaceName: req.workspaceName,
+		agentName:     req.agentName,
+		templateName:  req.templateName,
+		metricName:    m.Name,
+		labelsStr:     strings.Join(labelPairs, ","),
+	}
+}
+
+var _ prometheus.Collector = new(MetricsAggregator)
+
 func (am *annotatedMetric) asPrometheus() (prometheus.Metric, error) {
 	labels := make([]string, 0, len(agentMetricsLabels)+len(am.Labels))
 	labelValues := make([]string, 0, len(agentMetricsLabels)+len(am.Labels))
@@ -101,14 +127,25 @@ func NewMetricsAggregator(logger slog.Logger, registerer prometheus.Registerer,
 		metricsCleanupInterval = duration
 	}
 
+	storeSizeGauge := prometheus.NewGauge(prometheus.GaugeOpts{
+		Namespace: "coderd",
+		Subsystem: "prometheusmetrics",
+		Name:      "metrics_aggregator_store_size",
+		Help:      "The number of metrics stored in the aggregator",
+	})
+	err := registerer.Register(storeSizeGauge)
+	if err != nil {
+		return nil, err
+	}
+
 	updateHistogram := prometheus.NewHistogram(prometheus.HistogramOpts{
 		Namespace: "coderd",
 		Subsystem: "prometheusmetrics",
 		Name:      "metrics_aggregator_execution_update_seconds",
 		Help:      "Histogram for duration of metrics aggregator update in seconds.",
 		Buckets:   []float64{0.001, 0.005, 0.010, 0.025, 0.050, 0.100, 0.500, 1, 5, 10, 30},
 	})
-	err := registerer.Register(updateHistogram)
+	err = registerer.Register(updateHistogram)
 	if err != nil {
 		return nil, err
 	}
@@ -129,9 +166,12 @@ func NewMetricsAggregator(logger slog.Logger, registerer prometheus.Registerer,
 		log:                    logger.Named(loggerName),
 		metricsCleanupInterval: metricsCleanupInterval,
 
+		store: map[metricKey]annotatedMetric{},
+
 		collectCh: make(chan (chan []prometheus.Metric), sizeCollectCh),
 		updateCh:  make(chan updateRequest, sizeUpdateCh),
 
+		storeSizeGauge:   storeSizeGauge,
 		updateHistogram:  updateHistogram,
 		cleanupHistogram: cleanupHistogram,
 	}, nil
@@ -152,32 +192,32 @@ func (ma *MetricsAggregator) Run(ctx context.Context) func() {
 				ma.log.Debug(ctx, "update metrics")
 
 				timer := prometheus.NewTimer(ma.updateHistogram)
-			UpdateLoop:
 				for _, m := range req.metrics {
-					for i, q := range ma.queue {
-						if q.is(req, m) {
-							ma.queue[i].Stats_Metric.Value = m.Value
-							ma.queue[i].expiryDate = req.timestamp.Add(ma.metricsCleanupInterval)
-							continue UpdateLoop
+					key := hashKey(&req, m)
+
+					if val, ok := ma.store[key]; ok {
+						val.Stats_Metric.Value = m.Value
+						val.expiryDate = req.timestamp.Add(ma.metricsCleanupInterval)
+						ma.store[key] = val
+					} else {
+						ma.store[key] = annotatedMetric{
+							Stats_Metric:  m,
+							username:      req.username,
+							workspaceName: req.workspaceName,
+							agentName:     req.agentName,
+							templateName:  req.templateName,
+							expiryDate:    req.timestamp.Add(ma.metricsCleanupInterval),
 						}
 					}
-
-					ma.queue = append(ma.queue, annotatedMetric{
-						Stats_Metric:  m,
-						username:      req.username,
-						workspaceName: req.workspaceName,
-						agentName:     req.agentName,
-						templateName:  req.templateName,
-						expiryDate:    req.timestamp.Add(ma.metricsCleanupInterval),
-					})
 				}
-
 				timer.ObserveDuration()
+
+				ma.storeSizeGauge.Set(float64(len(ma.store)))
 			case outputCh := <-ma.collectCh:
 				ma.log.Debug(ctx, "collect metrics")
 
-				output := make([]prometheus.Metric, 0, len(ma.queue))
-				for _, m := range ma.queue {
+				output := make([]prometheus.Metric, 0, len(ma.store))
+				for _, m := range ma.store {
 					promMetric, err := m.asPrometheus()
 					if err != nil {
 						ma.log.Error(ctx, "can't convert Prometheus value type", slog.F("name", m.Name), slog.F("type", m.Type), slog.F("value", m.Value), slog.Error(err))
@@ -191,29 +231,17 @@ func (ma *MetricsAggregator) Run(ctx context.Context) func() {
 				ma.log.Debug(ctx, "clean expired metrics")
 
 				timer := prometheus.NewTimer(ma.cleanupHistogram)
-
 				now := time.Now()
 
-				var hasExpiredMetrics bool
-				for _, m := range ma.queue {
-					if now.After(m.expiryDate) {
-						hasExpiredMetrics = true
-						break
-					}
-				}
-
-				if hasExpiredMetrics {
-					fresh := make([]annotatedMetric, 0, len(ma.queue))
-					for _, m := range ma.queue {
-						if m.expiryDate.After(now) {
-							fresh = append(fresh, m)
-						}
+				for key, val := range ma.store {
+					if now.After(val.expiryDate) {
+						delete(ma.store, key)
 					}
-					ma.queue = fresh
 				}
 
 				timer.ObserveDuration()
 				cleanupTicker.Reset(ma.metricsCleanupInterval)
+				ma.storeSizeGauge.Set(float64(len(ma.store)))
 
 			case <-ctx.Done():
 				ma.log.Debug(ctx, "metrics aggregator is stopped")