perf: improve performance of metricsAggregator path by reducing memory allocations (#20724)

Signed-off-by: Callum Styan <callumstyan@gmail.com>
This commit is contained in:
Callum Styan
2025-11-24 15:45:08 -08:00
committed by GitHub
parent ca560d36ce
commit 658e8c34a9
3 changed files with 240 additions and 37 deletions
+100 -37
View File
@@ -37,6 +37,11 @@ const (
var MetricLabelValueEncoder = strings.NewReplacer("\\", "\\\\", "|", "\\|", ",", "\\,", "=", "\\=")
type descCacheEntry struct {
desc *prometheus.Desc
lastUsed time.Time
}
type MetricsAggregator struct {
store map[metricKey]annotatedMetric
@@ -50,6 +55,8 @@ type MetricsAggregator struct {
updateHistogram prometheus.Histogram
cleanupHistogram prometheus.Histogram
aggregateByLabels []string
// per-aggregator cache of descriptors
descCache map[string]descCacheEntry
}
type updateRequest struct {
@@ -107,42 +114,6 @@ func hashKey(req *updateRequest, m *agentproto.Stats_Metric) metricKey {
var _ prometheus.Collector = new(MetricsAggregator)
func (am *annotatedMetric) asPrometheus() (prometheus.Metric, error) {
var (
baseLabelNames = am.aggregateByLabels
baseLabelValues []string
extraLabels = am.Labels
)
for _, label := range baseLabelNames {
val, err := am.getFieldByLabel(label)
if err != nil {
return nil, err
}
baseLabelValues = append(baseLabelValues, val)
}
labels := make([]string, 0, len(baseLabelNames)+len(extraLabels))
labelValues := make([]string, 0, len(baseLabelNames)+len(extraLabels))
labels = append(labels, baseLabelNames...)
labelValues = append(labelValues, baseLabelValues...)
for _, l := range extraLabels {
labels = append(labels, l.Name)
labelValues = append(labelValues, l.Value)
}
desc := prometheus.NewDesc(am.Name, metricHelpForAgent, labels, nil)
valueType, err := asPrometheusValueType(am.Type)
if err != nil {
return nil, err
}
return prometheus.MustNewConstMetric(desc, valueType, am.Value, labelValues...), nil
}
// getFieldByLabel returns the related field value for a given label
func (am *annotatedMetric) getFieldByLabel(label string) (string, error) {
var labelVal string
@@ -364,7 +335,7 @@ func (ma *MetricsAggregator) Run(ctx context.Context) func() {
}
for _, m := range input {
promMetric, err := m.asPrometheus()
promMetric, err := ma.asPrometheus(&m)
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))
continue
@@ -386,6 +357,8 @@ func (ma *MetricsAggregator) Run(ctx context.Context) func() {
}
}
ma.cleanupDescCache()
timer.ObserveDuration()
cleanupTicker.Reset(ma.metricsCleanupInterval)
ma.storeSizeGauge.Set(float64(len(ma.store)))
@@ -407,6 +380,86 @@ func (ma *MetricsAggregator) Run(ctx context.Context) func() {
func (*MetricsAggregator) Describe(_ chan<- *prometheus.Desc) {
}
// cacheKeyForDesc is used to determine the cache key for a set of labels/extra labels. Used with the aggregators description cache.
// for strings.Builder returned errors from these functions are always nil.
// nolint:revive
func cacheKeyForDesc(name string, baseLabelNames []string, extraLabels []*agentproto.Stats_Metric_Label) string {
var b strings.Builder
hint := len(name) + (len(baseLabelNames)+len(extraLabels))*8
b.Grow(hint)
b.WriteString(name)
for _, ln := range baseLabelNames {
b.WriteByte('|')
b.WriteString(ln)
}
for _, l := range extraLabels {
b.WriteByte('|')
b.WriteString(l.Name)
}
return b.String()
}
// getOrCreateDec checks if we already have a metric description in the aggregators cache for a given combination of base
// labels and extra labels. If we do not, we create a new description and cache it.
func (ma *MetricsAggregator) getOrCreateDesc(name string, help string, baseLabelNames []string, extraLabels []*agentproto.Stats_Metric_Label) *prometheus.Desc {
if ma.descCache == nil {
ma.descCache = make(map[string]descCacheEntry)
}
key := cacheKeyForDesc(name, baseLabelNames, extraLabels)
if d, ok := ma.descCache[key]; ok {
d.lastUsed = time.Now()
ma.descCache[key] = d
return d.desc
}
nBase := len(baseLabelNames)
nExtra := len(extraLabels)
labels := make([]string, nBase+nExtra)
copy(labels, baseLabelNames)
for i, l := range extraLabels {
labels[nBase+i] = l.Name
}
d := prometheus.NewDesc(name, help, labels, nil)
ma.descCache[key] = descCacheEntry{d, time.Now()}
return d
}
// asPrometheus returns the annotatedMetric as a prometheus.Metric, it preallocates/fills by index, uses the aggregators
// metric description cache, and a small stack buffer for values in order to reduce memory allocations.
func (ma *MetricsAggregator) asPrometheus(am *annotatedMetric) (prometheus.Metric, error) {
baseLabelNames := am.aggregateByLabels
extraLabels := am.Labels
nBase := len(baseLabelNames)
nExtra := len(extraLabels)
nTotal := nBase + nExtra
var scratch [16]string
var labelValues []string
if nTotal <= len(scratch) {
labelValues = scratch[:nTotal]
} else {
labelValues = make([]string, nTotal)
}
for i, label := range baseLabelNames {
val, err := am.getFieldByLabel(label)
if err != nil {
return nil, err
}
labelValues[i] = val
}
for i, l := range extraLabels {
labelValues[nBase+i] = l.Value
}
desc := ma.getOrCreateDesc(am.Name, metricHelpForAgent, baseLabelNames, extraLabels)
valueType, err := asPrometheusValueType(am.Type)
if err != nil {
return nil, err
}
return prometheus.MustNewConstMetric(desc, valueType, am.Value, labelValues...), nil
}
var defaultAgentMetricsLabels = []string{agentmetrics.LabelUsername, agentmetrics.LabelWorkspaceName, agentmetrics.LabelAgentName, agentmetrics.LabelTemplateName}
// AgentMetricLabels are the labels used to decorate an agent's metrics.
@@ -453,6 +506,16 @@ func (ma *MetricsAggregator) Update(ctx context.Context, labels AgentMetricLabel
}
}
// Move to a function for testability
func (ma *MetricsAggregator) cleanupDescCache() {
now := time.Now()
for key, entry := range ma.descCache {
if now.Sub(entry.lastUsed) > ma.metricsCleanupInterval {
delete(ma.descCache, key)
}
}
}
func asPrometheusValueType(metricType agentproto.Stats_Metric_Type) (prometheus.ValueType, error) {
switch metricType {
case agentproto.Stats_Metric_GAUGE:
@@ -0,0 +1,89 @@
package prometheusmetrics
import (
"testing"
"time"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/require"
"cdr.dev/slog/sloggers/slogtest"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/agentmetrics"
"github.com/coder/coder/v2/testutil"
)
func TestDescCache_DescExpire(t *testing.T) {
const (
testWorkspaceName = "yogi-workspace"
testUsername = "yogi-bear"
testAgentName = "main-agent"
testTemplateName = "main-template"
)
testLabels := AgentMetricLabels{
Username: testUsername,
WorkspaceName: testWorkspaceName,
AgentName: testAgentName,
TemplateName: testTemplateName,
}
t.Parallel()
// given
registry := prometheus.NewRegistry()
ma, err := NewMetricsAggregator(slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), registry, time.Millisecond, agentmetrics.LabelAll)
require.NoError(t, err)
given := []*agentproto.Stats_Metric{
{Name: "a_counter_one", Type: agentproto.Stats_Metric_COUNTER, Value: 1},
}
_, err = ma.asPrometheus(&annotatedMetric{
given[0],
testLabels.Username,
testLabels.WorkspaceName,
testLabels.AgentName,
testLabels.TemplateName,
// the rest doesn't matter for this test
time.Now(),
[]string{},
})
require.NoError(t, err)
require.Eventually(t, func() bool {
ma.cleanupDescCache()
return len(ma.descCache) == 0
}, testutil.WaitShort, testutil.IntervalFast)
}
// TestDescCacheTimestampUpdate ensures that the timestamp update in getOrCreateDesc
// updates the map entry because d is a copy, not a pointer.
func TestDescCacheTimestampUpdate(t *testing.T) {
t.Parallel()
registry := prometheus.NewRegistry()
ma, err := NewMetricsAggregator(slogtest.Make(t, nil), registry, time.Hour, nil)
require.NoError(t, err)
baseLabelNames := []string{"label1", "label2"}
extraLabels := []*agentproto.Stats_Metric_Label{
{Name: "extra1", Value: "value1"},
}
desc1 := ma.getOrCreateDesc("test_metric", "help text", baseLabelNames, extraLabels)
require.NotNil(t, desc1)
key := cacheKeyForDesc("test_metric", baseLabelNames, extraLabels)
initialEntry := ma.descCache[key]
initialTime := initialEntry.lastUsed
desc2 := ma.getOrCreateDesc("test_metric", "help text", baseLabelNames, extraLabels)
require.NotNil(t, desc2)
updatedEntry := ma.descCache[key]
updatedTime := updatedEntry.lastUsed
require.NotEqual(t, initialTime, updatedTime,
"Timestamp was NOT updated in map when accessing a metric description that should be cached")
}
@@ -1,10 +1,12 @@
package prometheusmetrics
import (
"fmt"
"testing"
"github.com/stretchr/testify/require"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/agentmetrics"
)
@@ -36,3 +38,52 @@ func TestFilterAcceptableAgentLabels(t *testing.T) {
})
}
}
func benchAsPrometheus(b *testing.B, base []string, extraN int) {
am := annotatedMetric{
Stats_Metric: &agentproto.Stats_Metric{
Name: "blink_test_metric",
Type: agentproto.Stats_Metric_GAUGE,
Value: 1,
Labels: make([]*agentproto.Stats_Metric_Label, extraN),
},
username: "user",
workspaceName: "ws",
agentName: "agent",
templateName: "tmpl",
aggregateByLabels: base,
}
for i := 0; i < extraN; i++ {
am.Labels[i] = &agentproto.Stats_Metric_Label{Name: fmt.Sprintf("l%d", i), Value: "v"}
}
ma := &MetricsAggregator{}
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, err := ma.asPrometheus(&am)
if err != nil {
b.Fatal(err)
}
}
}
func Benchmark_asPrometheus(b *testing.B) {
cases := []struct {
name string
base []string
extraN int
}{
{"base4_extra0", defaultAgentMetricsLabels, 0},
{"base4_extra2", defaultAgentMetricsLabels, 2},
{"base4_extra5", defaultAgentMetricsLabels, 5},
{"base4_extra10", defaultAgentMetricsLabels, 10},
{"base2_extra5", []string{agentmetrics.LabelUsername, agentmetrics.LabelWorkspaceName}, 5},
}
for _, tc := range cases {
b.Run(tc.name, func(b *testing.B) {
benchAsPrometheus(b, tc.base, tc.extraN)
})
}
}