From dcecc600a11239664a4c295c08a283298481fc63 Mon Sep 17 00:00:00 2001
From: tdakkota <tanc13@yandex.ru>
Date: Tue, 12 Mar 2024 02:18:53 +0300
Subject: [PATCH] test(prompb): ensure that slicepool is properly re-used

---
 internal/prompb/alloc_test.go | 73 +++++++++++++++++++++++++++++++++++
 internal/prompb/pool.go       | 36 +++++++++++++++++
 internal/prompb/remote.pb.go  | 15 +------
 3 files changed, 111 insertions(+), 13 deletions(-)
 create mode 100644 internal/prompb/alloc_test.go

diff --git a/internal/prompb/alloc_test.go b/internal/prompb/alloc_test.go
new file mode 100644
index 00000000..a724905b
--- /dev/null
+++ b/internal/prompb/alloc_test.go
@@ -0,0 +1,73 @@
+package prompb
+
+import (
+	"testing"
+
+	nativeprompb "github.com/prometheus/prometheus/prompb"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEnsurePoolReuse(t *testing.T) {
+	labels := []nativeprompb.Label{
+		{Name: "a", Value: "1"},
+		{Name: "b", Value: "2"},
+		{Name: "c", Value: "3"},
+		{Name: "d", Value: "4"},
+	}
+	samples := []nativeprompb.Sample{
+		{Value: 1, Timestamp: 1},
+		{Value: 1, Timestamp: 2},
+		{Value: 1, Timestamp: 3},
+		{Value: 1, Timestamp: 4},
+	}
+
+	data1, err := (&nativeprompb.WriteRequest{
+		Timeseries: []nativeprompb.TimeSeries{
+			{
+				Labels: labels[:2],
+			},
+			{
+				Samples: samples[:4],
+			},
+		},
+	}).Marshal()
+	require.NoError(t, err)
+
+	data2, err := (&nativeprompb.WriteRequest{
+		Timeseries: []nativeprompb.TimeSeries{
+			{
+				Samples: samples[:2],
+			},
+			{
+				Labels: labels[:4],
+			},
+		},
+	}).Marshal()
+	require.NoError(t, err)
+
+	// Fill target pool.
+	target := WriteRequest{
+		Timeseries: make([]TimeSeries, 0, 2),
+		Pools: &Pools{
+			Labels: &slicepool[Label]{
+				pool: make([]Label, 0, 4),
+			},
+			Samples: &slicepool[Sample]{
+				pool: make([]Sample, 0, 4),
+			},
+		},
+	}
+	target.Pools.init()
+
+	// Ensure that slicepool is properly re-used.
+	allocs := testing.AllocsPerRun(10, func() {
+		if err := target.Unmarshal(data1); err != nil {
+			t.Fatal(err)
+		}
+		target.Reset()
+		if err := target.Unmarshal(data2); err != nil {
+			t.Fatal(err)
+		}
+	})
+	require.Zero(t, allocs)
+}
diff --git a/internal/prompb/pool.go b/internal/prompb/pool.go
index fe416a03..cb30c284 100644
--- a/internal/prompb/pool.go
+++ b/internal/prompb/pool.go
@@ -29,6 +29,42 @@ type Pools struct {
 	HistogramPositiveCounts *slicepool[float64]
 }
 
+func (p *Pools) init() {
+	if p.Labels == nil {
+		p.Labels = new(slicepool[Label])
+	}
+	if p.Samples == nil {
+		p.Samples = new(slicepool[Sample])
+	}
+	if p.Exemplars == nil {
+		p.Exemplars = new(slicepool[Exemplar])
+	}
+	if p.ExemplarLabels == nil {
+		p.ExemplarLabels = new(slicepool[Label])
+	}
+	if p.Histograms == nil {
+		p.Histograms = new(slicepool[Histogram])
+	}
+	if p.HistogramNegativeSpans == nil {
+		p.HistogramNegativeSpans = new(slicepool[BucketSpan])
+	}
+	if p.HistogramNegativeDeltas == nil {
+		p.HistogramNegativeDeltas = new(slicepool[int64])
+	}
+	if p.HistogramNegativeCounts == nil {
+		p.HistogramNegativeCounts = new(slicepool[float64])
+	}
+	if p.HistogramPositiveSpans == nil {
+		p.HistogramPositiveSpans = new(slicepool[BucketSpan])
+	}
+	if p.HistogramPositiveDeltas == nil {
+		p.HistogramPositiveDeltas = new(slicepool[int64])
+	}
+	if p.HistogramPositiveCounts == nil {
+		p.HistogramPositiveCounts = new(slicepool[float64])
+	}
+}
+
 func (p *Pools) Reset() {
 	if p == nil {
 		return
diff --git a/internal/prompb/remote.pb.go b/internal/prompb/remote.pb.go
index 07c7a00e..76367cf3 100644
--- a/internal/prompb/remote.pb.go
+++ b/internal/prompb/remote.pb.go
@@ -14,19 +14,8 @@ type WriteRequest struct {
 // Unmarshal unmarshals WriteRequest from src.
 func (req *WriteRequest) Unmarshal(src []byte) (err error) {
 	if req.Pools == nil {
-		req.Pools = &Pools{
-			Labels:                  new(slicepool[Label]),
-			Samples:                 new(slicepool[Sample]),
-			Exemplars:               new(slicepool[Exemplar]),
-			ExemplarLabels:          new(slicepool[Label]),
-			Histograms:              new(slicepool[Histogram]),
-			HistogramNegativeSpans:  new(slicepool[BucketSpan]),
-			HistogramNegativeDeltas: new(slicepool[int64]),
-			HistogramNegativeCounts: new(slicepool[float64]),
-			HistogramPositiveSpans:  new(slicepool[BucketSpan]),
-			HistogramPositiveDeltas: new(slicepool[int64]),
-			HistogramPositiveCounts: new(slicepool[float64]),
-		}
+		req.Pools = &Pools{}
+		req.Pools.init()
 	}
 
 	var fc easyproto.FieldContext