From a09a1d34cbc74daa8ed70234b99467a30b020a40 Mon Sep 17 00:00:00 2001
From: Bulat Khasanov <afti@yandex.ru>
Date: Wed, 7 Jun 2023 10:39:02 +0300
Subject: [PATCH] Reduce constrainLabels allocations (#1272)

* Add bench

Signed-off-by: Bulat Khasanov <afti@yandex.ru>

* Reduce constrainLabels allocations

Signed-off-by: Bulat Khasanov <afti@yandex.ru>

---------

Signed-off-by: Bulat Khasanov <afti@yandex.ru>
---
 prometheus/vec.go      | 35 ++++++++++++++++++++++++++++++-----
 prometheus/vec_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/prometheus/vec.go b/prometheus/vec.go
index 386fb2d..f0d0015 100644
--- a/prometheus/vec.go
+++ b/prometheus/vec.go
@@ -20,6 +20,24 @@ import (
 	"github.com/prometheus/common/model"
 )
 
+var labelsPool = &sync.Pool{
+	New: func() interface{} {
+		return make(Labels)
+	},
+}
+
+func getLabelsFromPool() Labels {
+	return labelsPool.Get().(Labels)
+}
+
+func putLabelsToPool(labels Labels) {
+	for k := range labels {
+		delete(labels, k)
+	}
+
+	labelsPool.Put(labels)
+}
+
 // MetricVec is a Collector to bundle metrics of the same name that differ in
 // their label values. MetricVec is not used directly but as a building block
 // for implementations of vectors of a given metric type, like GaugeVec,
@@ -93,6 +111,8 @@ func (m *MetricVec) DeleteLabelValues(lvs ...string) bool {
 // there for pros and cons of the two methods.
 func (m *MetricVec) Delete(labels Labels) bool {
 	labels = constrainLabels(m.desc, labels)
+	defer putLabelsToPool(labels)
+
 	h, err := m.hashLabels(labels)
 	if err != nil {
 		return false
@@ -109,6 +129,8 @@ func (m *MetricVec) Delete(labels Labels) bool {
 // To match curried labels with DeletePartialMatch, it must be called on the base vector.
 func (m *MetricVec) DeletePartialMatch(labels Labels) int {
 	labels = constrainLabels(m.desc, labels)
+	defer putLabelsToPool(labels)
+
 	return m.metricMap.deleteByLabels(labels, m.curry)
 }
 
@@ -229,6 +251,8 @@ func (m *MetricVec) GetMetricWithLabelValues(lvs ...string) (Metric, error) {
 // for example GaugeVec.
 func (m *MetricVec) GetMetricWith(labels Labels) (Metric, error) {
 	labels = constrainLabels(m.desc, labels)
+	defer putLabelsToPool(labels)
+
 	h, err := m.hashLabels(labels)
 	if err != nil {
 		return nil, err
@@ -647,15 +671,16 @@ func inlineLabelValues(lvs []string, curry []curriedLabelValue) []string {
 }
 
 func constrainLabels(desc *Desc, labels Labels) Labels {
-	constrainedValues := make(Labels, len(labels))
+	constrainedLabels := getLabelsFromPool()
 	for l, v := range labels {
 		if i, ok := indexOf(l, desc.variableLabels.labelNames()); ok {
-			constrainedValues[l] = desc.variableLabels[i].Constrain(v)
-			continue
+			v = desc.variableLabels[i].Constrain(v)
 		}
-		constrainedValues[l] = v
+
+		constrainedLabels[l] = v
 	}
-	return constrainedValues
+
+	return constrainedLabels
 }
 
 func constrainLabelValues(desc *Desc, lvs []string, curry []curriedLabelValue) []string {
diff --git a/prometheus/vec_test.go b/prometheus/vec_test.go
index 63f52af..89bd854 100644
--- a/prometheus/vec_test.go
+++ b/prometheus/vec_test.go
@@ -904,6 +904,13 @@ func testConstrainedCurryVec(t *testing.T, vec *CounterVec, constraint func(stri
 	})
 }
 
+func BenchmarkMetricVecWithBasic(b *testing.B) {
+	benchmarkMetricVecWith(b, Labels{
+		"l1": "onevalue",
+		"l2": "twovalue",
+	})
+}
+
 func BenchmarkMetricVecWithLabelValuesBasic(b *testing.B) {
 	benchmarkMetricVecWithLabelValues(b, map[string][]string{
 		"l1": {"onevalue"},
@@ -948,6 +955,27 @@ func benchmarkMetricVecWithLabelValuesCardinality(b *testing.B, nkeys, nvalues i
 	benchmarkMetricVecWithLabelValues(b, labels)
 }
 
+func benchmarkMetricVecWith(b *testing.B, labels map[string]string) {
+	var keys []string
+	for k := range labels {
+		keys = append(keys, k)
+	}
+
+	vec := NewGaugeVec(
+		GaugeOpts{
+			Name: "test",
+			Help: "helpless",
+		},
+		keys,
+	)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		vec.With(labels)
+	}
+}
+
 func benchmarkMetricVecWithLabelValues(b *testing.B, labels map[string][]string) {
 	var keys []string
 	for k := range labels { // Map order dependent, who cares though.