prometheusremotewriteexporter: Remove unnecessary buffer copy in proto conversion

Signed-off-by: Bogdan Drutu <[email protected]>

Author: bogdandrutu

exporter/prometheusremotewriteexporter/exporter.go

@@ -17,8 +17,7 @@ import (
 	"time"
 
 	"github.com/cenkalti/backoff/v5"
-	"github.com/gogo/protobuf/proto"
-	"github.com/golang/snappy"
+	"github.com/klauspost/compress/snappy"
 	"github.com/prometheus/otlptranslator"
 	"github.com/prometheus/prometheus/config"
 	"github.com/prometheus/prometheus/prompb"
@@ -80,16 +79,46 @@ func (p *prwTelemetryOtel) recordWrittenExemplars(ctx context.Context, numExempl
 	p.telemetryBuilder.ExporterPrometheusremotewriteWrittenExemplars.Add(ctx, numExemplars, metric.WithAttributes(p.otelAttrs...))
 }
 
+type gogoProto interface {
+	Size() int
+	MarshalToSizedBuffer([]byte) (int, error)
+}
+
 type buffer struct {
-	protobuf *proto.Buffer
+	protobuf []byte
 	snappy   []byte
 }
 
+func (b *buffer) MarshalAndEncode(req gogoProto) ([]byte, error) {
+	sizePb := req.Size()
+	if sizePb > cap(b.protobuf) {
+		b.protobuf = make([]byte, sizePb)
+	}
+	b.protobuf = b.protobuf[:sizePb]
+	n, err := req.MarshalToSizedBuffer(b.protobuf)
+	if err != nil {
+		return nil, err
+	}
+	b.protobuf = b.protobuf[:n]
+
+	// If we don't pass a buffer large enough, Snappy Encode function will not use it and instead will allocate a new buffer.
+	// Manually grow the buffer to make sure Snappy uses it and we can re-use it afterwards.
+	maxCompressedLen := snappy.MaxEncodedLen(len(b.protobuf))
+	if maxCompressedLen > len(b.snappy) {
+		if cap(b.snappy) < maxCompressedLen {
+			b.snappy = make([]byte, maxCompressedLen)
+		} else {
+			b.snappy = b.snappy[:maxCompressedLen]
+		}
+	}
+	return snappy.Encode(b.snappy, b.protobuf), nil
+}
+
 // A reusable buffer pool for serializing protobufs and compressing them with Snappy.
 var bufferPool = sync.Pool{
 	New: func() any {
 		return &buffer{
-			protobuf: proto.NewBuffer(nil),
+			protobuf: nil,
 			snappy:   nil,
 		}
 	},
@@ -340,6 +369,8 @@ func (prwe *prwExporter) export(ctx context.Context, requests []*prompb.WriteReq
 	for i := 0; i < concurrencyLimit; i++ {
 		go func() {
 			defer wg.Done()
+			buf := bufferPool.Get().(*buffer)
+			defer bufferPool.Put(buf)
 			for {
 				select {
 				case <-ctx.Done(): // Check firstly to ensure that the context wasn't cancelled.
@@ -350,18 +381,15 @@ func (prwe *prwExporter) export(ctx context.Context, requests []*prompb.WriteReq
 						return
 					}
 
-					buf := bufferPool.Get().(*buffer)
-					buf.protobuf.Reset()
-					defer bufferPool.Put(buf)
-
-					if errMarshal := buf.protobuf.Marshal(request); errMarshal != nil {
+					reqBuf, errMarshal := buf.MarshalAndEncode(request)
+					if errMarshal != nil {
 						mu.Lock()
 						errs = multierr.Append(errs, consumererror.NewPermanent(errMarshal))
 						mu.Unlock()
 						return
 					}
 
-					if errExecute := prwe.execute(ctx, buf); errExecute != nil {
+					if errExecute := prwe.execute(ctx, reqBuf); errExecute != nil {
 						mu.Lock()
 						errs = multierr.Append(errs, consumererror.NewPermanent(errExecute))
 						mu.Unlock()
@@ -375,19 +403,7 @@ func (prwe *prwExporter) export(ctx context.Context, requests []*prompb.WriteReq
 	return errs
 }
 
-func (prwe *prwExporter) execute(ctx context.Context, buf *buffer) error {
-	// If we don't pass a buffer large enough, Snappy Encode function will not use it and instead will allocate a new buffer.
-	// Manually grow the buffer to make sure Snappy uses it and we can re-use it afterwards.
-	maxCompressedLen := snappy.MaxEncodedLen(len(buf.protobuf.Bytes()))
-	if maxCompressedLen > len(buf.snappy) {
-		if cap(buf.snappy) < maxCompressedLen {
-			buf.snappy = make([]byte, maxCompressedLen)
-		} else {
-			buf.snappy = buf.snappy[:maxCompressedLen]
-		}
-	}
-	compressedData := snappy.Encode(buf.snappy, buf.protobuf.Bytes())
-
+func (prwe *prwExporter) execute(ctx context.Context, buf []byte) error {
 	retryCount := 0
 	// executeFunc can be used for backoff and non backoff scenarios.
 	executeFunc := func() (int, error) {
@@ -402,7 +418,7 @@ func (prwe *prwExporter) execute(ctx context.Context, buf *buffer) error {
 		}
 
 		// Create the HTTP POST request to send to the endpoint
-		req, err := http.NewRequestWithContext(ctx, http.MethodPost, prwe.endpointURL.String(), bytes.NewReader(compressedData))
+		req, err := http.NewRequestWithContext(ctx, http.MethodPost, prwe.endpointURL.String(), bytes.NewReader(buf))
 		if err != nil {
 			return http.StatusBadRequest, backoff.Permanent(consumererror.NewPermanent(err))
 		}

exporter/prometheusremotewriteexporter/exporter_test.go

@@ -1195,16 +1195,15 @@ func TestRetries(t *testing.T) {
 				telemetry: telemetry,
 			}
 			buf := bufferPool.Get().(*buffer)
-			buf.protobuf.Reset()
 			defer bufferPool.Put(buf)
 
-			errMarshal := buf.protobuf.Marshal(&prompb.WriteRequest{})
+			reqBuf, errMarshal := buf.MarshalAndEncode(&prompb.WriteRequest{})
 			if errMarshal != nil {
 				require.NoError(t, errMarshal)
 				return
 			}
 
-			err = exporter.execute(tt.ctx, buf)
+			err = exporter.execute(tt.ctx, reqBuf)
 			tt.assertError(t, err)
 			tt.assertErrorType(t, err)
 			assert.Equal(t, tt.expectedAttempts, totalAttempts)
@@ -1228,10 +1227,17 @@ func benchmarkExecute(b *testing.B, numSample int) {
 	endpointURL, err := url.Parse(mockServer.URL)
 	require.NoError(b, err)
 
+	// Create the telemetry
+	testTel := componenttest.NewTelemetry()
+	telemetry, err := newPRWTelemetry(exporter.Settings{TelemetrySettings: testTel.NewTelemetrySettings()}, endpointURL)
+	require.NoError(b, err)
+
 	// Create the prwExporter
 	exporter := &prwExporter{
 		endpointURL: endpointURL,
 		client:      http.DefaultClient,
+		settings:    testTel.NewTelemetrySettings(),
+		telemetry:   telemetry,
 	}
 
 	generateSamples := func(n int) []prompb.Sample {
@@ -1297,18 +1303,18 @@ func benchmarkExecute(b *testing.B, numSample int) {
 	ctx := b.Context()
 	b.ReportAllocs()
 	b.ResetTimer()
+
 	for _, req := range reqs {
 		buf := bufferPool.Get().(*buffer)
-		buf.protobuf.Reset()
-		defer bufferPool.Put(buf)
-
-		errMarshal := buf.protobuf.Marshal(req)
+		reqBuf, errMarshal := buf.MarshalAndEncode(req)
 		if errMarshal != nil {
 			require.NoError(b, errMarshal)
 			return
 		}
-		err := exporter.execute(ctx, buf)
-		require.NoError(b, err)
+		if err = exporter.execute(ctx, reqBuf); err != nil {
+			b.Fatal(err)
+		}
+		bufferPool.Put(buf)
 	}
 }
 

exporter/prometheusremotewriteexporter/exporter_v2.go

@@ -51,6 +51,8 @@ func (prwe *prwExporter) exportV2(ctx context.Context, requests []*writev2.Reque
 	for i := 0; i < concurrencyLimit; i++ {
 		go func() {
 			defer wg.Done()
+			buf := bufferPool.Get().(*buffer)
+			defer bufferPool.Put(buf)
 			for {
 				select {
 				case <-ctx.Done(): // Check firstly to ensure that the context wasn't cancelled.
@@ -61,10 +63,7 @@ func (prwe *prwExporter) exportV2(ctx context.Context, requests []*writev2.Reque
 						return
 					}
 
-					buf := bufferPool.Get().(*buffer)
-					buf.protobuf.Reset()
-
-					errMarshal := buf.protobuf.Marshal(request)
+					reqBuf, errMarshal := buf.MarshalAndEncode(request)
 					if errMarshal != nil {
 						mu.Lock()
 						errs = multierr.Append(errs, errMarshal)
@@ -73,7 +72,7 @@ func (prwe *prwExporter) exportV2(ctx context.Context, requests []*writev2.Reque
 						return
 					}
 
-					if errExecute := prwe.execute(ctx, buf); errExecute != nil {
+					if errExecute := prwe.execute(ctx, reqBuf); errExecute != nil {
 						mu.Lock()
 						errs = multierr.Append(errs, errExecute)
 						mu.Unlock()