[exporter/prometheusremotewrite] Fix data race in batch series state if called concurrently

.chloggen/prwexporter-batchseries-concurrency-bugfix.yaml

@@ -0,0 +1,26 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: bug_fix
+
+# The name of the component, or a single word describing the area of concern, (e.g. filelogreceiver)
+component: prometheusremotewriteexporter
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: "Fix data race in batch series state."
+
+# Mandatory: One or more tracking issues related to the change. You can use the PR number here if no issue exists.
+issues: [36524]
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext: ""
+# If your change doesn't affect end users or the exported elements of any package,
+# you should instead start your pull request title with [chore] or use the "Skip Changelog" label.
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: [user]

exporter/prometheusremotewriteexporter/exporter.go

@@ -84,7 +84,7 @@ type prwExporter struct {
 	wal                  *prweWAL
 	exporterSettings     prometheusremotewrite.Settings
 	telemetry            prwTelemetry
-	batchTimeSeriesState batchTimeSeriesState
+	batchTimeSeriesState *batchTimeSeriesState
 }
 
 func newPRWTelemetry(set exporter.Settings) (prwTelemetry, error) {
@@ -191,7 +191,6 @@ func (prwe *prwExporter) PushMetrics(ctx context.Context, md pmetric.Metrics) er
 	case <-prwe.closeChan:
 		return errors.New("shutdown has been called")
 	default:
-
 		tsMap, err := prometheusremotewrite.FromMetrics(md, prwe.exporterSettings)
 		if err != nil {
 			prwe.telemetry.recordTranslationFailure(ctx)
@@ -229,7 +228,7 @@ func (prwe *prwExporter) handleExport(ctx context.Context, tsMap map[string]*pro
 	}
 
 	// Calls the helper function to convert and batch the TsMap to the desired format
-	requests, err := batchTimeSeries(tsMap, prwe.maxBatchSizeBytes, m, &prwe.batchTimeSeriesState)
+	requests, err := batchTimeSeries(tsMap, prwe.maxBatchSizeBytes, m, prwe.batchTimeSeriesState)
 	if err != nil {
 		return err
 	}

exporter/prometheusremotewriteexporter/exporter_concurrency_test.go

@@ -0,0 +1,138 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package prometheusremotewriteexporter
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/golang/snappy"
+	"github.com/prometheus/prometheus/prompb"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/component/componenttest"
+	"go.opentelemetry.io/collector/config/confighttp"
+	"go.opentelemetry.io/collector/config/configretry"
+	"go.opentelemetry.io/collector/config/configtelemetry"
+	"go.opentelemetry.io/collector/exporter/exportertest"
+	"go.opentelemetry.io/collector/pdata/pmetric"
+
+	"github.com/open-telemetry/opentelemetry-collector-contrib/internal/coreinternal/testdata"
+)
+
+// Test everything works when there is more than one goroutine calling PushMetrics.
+// Today we only use 1 worker per exporter, but the intention of this test is to future-proof in case it changes.
+func Test_PushMetricsConcurrent(t *testing.T) {
+	n := 1000
+	ms := make([]pmetric.Metrics, n)
+	testIDKey := "test_id"
+	for i := 0; i < n; i++ {
+		m := testdata.GenerateMetricsOneMetric()
+		dps := m.ResourceMetrics().At(0).ScopeMetrics().At(0).Metrics().At(0).Sum().DataPoints()
+		for j := 0; j < dps.Len(); j++ {
+			dp := dps.At(j)
+			dp.Attributes().PutInt(testIDKey, int64(i))
+		}
+		ms[i] = m
+	}
+	received := make(map[int]prompb.TimeSeries)
+	var mu sync.Mutex
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		assert.NotNil(t, body)
+		// Receives the http requests and unzip, unmarshalls, and extracts TimeSeries
+		assert.Equal(t, "0.1.0", r.Header.Get("X-Prometheus-Remote-Write-Version"))
+		assert.Equal(t, "snappy", r.Header.Get("Content-Encoding"))
+		var unzipped []byte
+
+		dest, err := snappy.Decode(unzipped, body)
+		assert.NoError(t, err)
+
+		wr := &prompb.WriteRequest{}
+		ok := proto.Unmarshal(dest, wr)
+		assert.NoError(t, ok)
+		assert.Len(t, wr.Timeseries, 2)
+		ts := wr.Timeseries[0]
+		foundLabel := false
+		for _, label := range ts.Labels {
+			if label.Name == testIDKey {
+				id, err := strconv.Atoi(label.Value)
+				assert.NoError(t, err)
+				mu.Lock()
+				_, ok := received[id]
+				assert.False(t, ok) // fail if we already saw it
+				received[id] = ts
+				mu.Unlock()
+				foundLabel = true
+				break
+			}
+		}
+		assert.True(t, foundLabel)
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	defer server.Close()
+
+	// Adjusted retry settings for faster testing
+	retrySettings := configretry.BackOffConfig{
+		Enabled:         true,
+		InitialInterval: 100 * time.Millisecond, // Shorter initial interval
+		MaxInterval:     1 * time.Second,        // Shorter max interval
+		MaxElapsedTime:  2 * time.Second,        // Shorter max elapsed time
+	}
+	clientConfig := confighttp.NewDefaultClientConfig()
+	clientConfig.Endpoint = server.URL
+	clientConfig.ReadBufferSize = 0
+	clientConfig.WriteBufferSize = 512 * 1024
+	cfg := &Config{
+		Namespace:         "",
+		ClientConfig:      clientConfig,
+		MaxBatchSizeBytes: 3000000,
+		RemoteWriteQueue:  RemoteWriteQueue{NumConsumers: 1},
+		TargetInfo: &TargetInfo{
+			Enabled: true,
+		},
+		CreatedMetric: &CreatedMetric{
+			Enabled: false,
+		},
+		BackOffConfig: retrySettings,
+	}
+
+	assert.NotNil(t, cfg)
+	set := exportertest.NewNopSettings()
+	set.MetricsLevel = configtelemetry.LevelBasic
+
+	prwe, nErr := newPRWExporter(cfg, set)
+
+	require.NoError(t, nErr)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	require.NoError(t, prwe.Start(ctx, componenttest.NewNopHost()))
+	defer func() {
+		require.NoError(t, prwe.Shutdown(ctx))
+	}()
+
+	var wg sync.WaitGroup
+	wg.Add(n)
+	for _, m := range ms {
+		go func() {
+			err := prwe.PushMetrics(ctx, m)
+			assert.NoError(t, err)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	assert.Len(t, received, n)
+}

exporter/prometheusremotewriteexporter/helper.go

@@ -7,24 +7,29 @@ import (
 	"errors"
 	"math"
 	"sort"
+	"sync/atomic"
 
 	"github.com/prometheus/prometheus/prompb"
 )
 
 type batchTimeSeriesState struct {
 	// Track batch sizes sent to avoid over allocating huge buffers.
 	// This helps in the case where large batches are sent to avoid allocating too much unused memory
-	nextTimeSeriesBufferSize     int
-	nextMetricMetadataBufferSize int
-	nextRequestBufferSize        int
+	nextTimeSeriesBufferSize     atomic.Int64
+	nextMetricMetadataBufferSize atomic.Int64
+	nextRequestBufferSize        atomic.Int64
 }
 
-func newBatchTimeSericesState() batchTimeSeriesState {
-	return batchTimeSeriesState{
-		nextTimeSeriesBufferSize:     math.MaxInt,
-		nextMetricMetadataBufferSize: math.MaxInt,
-		nextRequestBufferSize:        0,
+func newBatchTimeSericesState() *batchTimeSeriesState {
+	state := &batchTimeSeriesState{
+		nextTimeSeriesBufferSize:     atomic.Int64{},
+		nextMetricMetadataBufferSize: atomic.Int64{},
+		nextRequestBufferSize:        atomic.Int64{},
 	}
+	state.nextTimeSeriesBufferSize.Store(math.MaxInt64)
+	state.nextMetricMetadataBufferSize.Store(math.MaxInt64)
+	state.nextRequestBufferSize.Store(0)
+	return state
 }
 
 // batchTimeSeries splits series into multiple batch write requests.
@@ -34,22 +39,22 @@ func batchTimeSeries(tsMap map[string]*prompb.TimeSeries, maxBatchByteSize int,
 	}
 
 	// Allocate a buffer size of at least 10, or twice the last # of requests we sent
-	requests := make([]*prompb.WriteRequest, 0, max(10, state.nextRequestBufferSize))
+	requests := make([]*prompb.WriteRequest, 0, max(10, state.nextRequestBufferSize.Load()))
 
 	// Allocate a time series buffer 2x the last time series batch size or the length of the input if smaller
-	tsArray := make([]prompb.TimeSeries, 0, min(state.nextTimeSeriesBufferSize, len(tsMap)))
+	tsArray := make([]prompb.TimeSeries, 0, min(state.nextTimeSeriesBufferSize.Load(), int64(len(tsMap))))
 	sizeOfCurrentBatch := 0
 
 	i := 0
 	for _, v := range tsMap {
 		sizeOfSeries := v.Size()
 
 		if sizeOfCurrentBatch+sizeOfSeries >= maxBatchByteSize {
-			state.nextTimeSeriesBufferSize = max(10, 2*len(tsArray))
+			state.nextTimeSeriesBufferSize.Store(int64(max(10, 2*len(tsArray))))
 			wrapped := convertTimeseriesToRequest(tsArray)
 			requests = append(requests, wrapped)
 
-			tsArray = make([]prompb.TimeSeries, 0, min(state.nextTimeSeriesBufferSize, len(tsMap)-i))
+			tsArray = make([]prompb.TimeSeries, 0, min(state.nextTimeSeriesBufferSize.Load(), int64(len(tsMap)-i)))
 			sizeOfCurrentBatch = 0
 		}
 
@@ -64,18 +69,18 @@ func batchTimeSeries(tsMap map[string]*prompb.TimeSeries, maxBatchByteSize int,
 	}
 
 	// Allocate a metric metadata buffer 2x the last metric metadata batch size or the length of the input if smaller
-	mArray := make([]prompb.MetricMetadata, 0, min(state.nextMetricMetadataBufferSize, len(m)))
+	mArray := make([]prompb.MetricMetadata, 0, min(state.nextMetricMetadataBufferSize.Load(), int64(len(m))))
 	sizeOfCurrentBatch = 0
 	i = 0
 	for _, v := range m {
 		sizeOfM := v.Size()
 
 		if sizeOfCurrentBatch+sizeOfM >= maxBatchByteSize {
-			state.nextMetricMetadataBufferSize = max(10, 2*len(mArray))
+			state.nextMetricMetadataBufferSize.Store(int64(max(10, 2*len(mArray))))
 			wrapped := convertMetadataToRequest(mArray)
 			requests = append(requests, wrapped)
 
-			mArray = make([]prompb.MetricMetadata, 0, min(state.nextMetricMetadataBufferSize, len(m)-i))
+			mArray = make([]prompb.MetricMetadata, 0, min(state.nextMetricMetadataBufferSize.Load(), int64(len(m)-i)))
 			sizeOfCurrentBatch = 0
 		}
 
@@ -89,7 +94,7 @@ func batchTimeSeries(tsMap map[string]*prompb.TimeSeries, maxBatchByteSize int,
 		requests = append(requests, wrapped)
 	}
 
-	state.nextRequestBufferSize = 2 * len(requests)
+	state.nextRequestBufferSize.Store(int64(2 * len(requests)))
 	return requests, nil
 }
 

exporter/prometheusremotewriteexporter/helper_test.go

@@ -59,7 +59,7 @@ func Test_batchTimeSeries(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			state := newBatchTimeSericesState()
-			requests, err := batchTimeSeries(tt.tsMap, tt.maxBatchByteSize, nil, &state)
+			requests, err := batchTimeSeries(tt.tsMap, tt.maxBatchByteSize, nil, state)
 			if tt.returnErr {
 				assert.Error(t, err)
 				return
@@ -68,13 +68,13 @@ func Test_batchTimeSeries(t *testing.T) {
 			assert.NoError(t, err)
 			assert.Len(t, requests, tt.numExpectedRequests)
 			if tt.numExpectedRequests <= 1 {
-				assert.Equal(t, math.MaxInt, state.nextTimeSeriesBufferSize)
-				assert.Equal(t, math.MaxInt, state.nextMetricMetadataBufferSize)
-				assert.Equal(t, 2*len(requests), state.nextRequestBufferSize)
+				assert.Equal(t, int64(math.MaxInt64), state.nextTimeSeriesBufferSize.Load())
+				assert.Equal(t, int64(math.MaxInt64), state.nextMetricMetadataBufferSize.Load())
+				assert.Equal(t, int64(2*len(requests)), state.nextRequestBufferSize.Load())
 			} else {
-				assert.Equal(t, max(10, len(requests[len(requests)-2].Timeseries)*2), state.nextTimeSeriesBufferSize)
-				assert.Equal(t, math.MaxInt, state.nextMetricMetadataBufferSize)
-				assert.Equal(t, 2*len(requests), state.nextRequestBufferSize)
+				assert.Equal(t, int64(max(10, len(requests[len(requests)-2].Timeseries)*2)), state.nextTimeSeriesBufferSize.Load())
+				assert.Equal(t, int64(math.MaxInt64), state.nextMetricMetadataBufferSize.Load())
+				assert.Equal(t, int64(2*len(requests)), state.nextRequestBufferSize.Load())
 			}
 		})
 	}
@@ -97,13 +97,13 @@ func Test_batchTimeSeriesUpdatesStateForLargeBatches(t *testing.T) {
 	tsMap1 := getTimeseriesMap(tsArray)
 
 	state := newBatchTimeSericesState()
-	requests, err := batchTimeSeries(tsMap1, 1000000, nil, &state)
+	requests, err := batchTimeSeries(tsMap1, 1000000, nil, state)
 
 	assert.NoError(t, err)
 	assert.Len(t, requests, 18)
-	assert.Equal(t, len(requests[len(requests)-2].Timeseries)*2, state.nextTimeSeriesBufferSize)
-	assert.Equal(t, math.MaxInt, state.nextMetricMetadataBufferSize)
-	assert.Equal(t, 36, state.nextRequestBufferSize)
+	assert.Equal(t, int64(len(requests[len(requests)-2].Timeseries)*2), state.nextTimeSeriesBufferSize.Load())
+	assert.Equal(t, int64(math.MaxInt64), state.nextMetricMetadataBufferSize.Load())
+	assert.Equal(t, int64(36), state.nextRequestBufferSize.Load())
 }
 
 // Benchmark_batchTimeSeries checks batchTimeSeries
@@ -132,7 +132,7 @@ func Benchmark_batchTimeSeries(b *testing.B) {
 	state := newBatchTimeSericesState()
 	// Run batchTimeSeries 100 times with a 1mb max request size
 	for i := 0; i < b.N; i++ {
-		requests, err := batchTimeSeries(tsMap1, 1000000, nil, &state)
+		requests, err := batchTimeSeries(tsMap1, 1000000, nil, state)
 		assert.NoError(b, err)
 		assert.Len(b, requests, 18)
 	}