Add batched exports to the Prometheus Remote Write Exporter

exporter/prometheusremotewriteexporter/exporter.go

@@ -143,9 +143,9 @@ func (prwe *PrwExporter) PushMetrics(ctx context.Context, md pdata.Metrics) (int
 			}
 		}
 
-		if err := prwe.export(ctx, tsMap); err != nil {
+		if exportErrors := prwe.export(ctx, tsMap); exportErrors != nil {
 			dropped = md.MetricCount()
-			errs = append(errs, consumererror.Permanent(err))
+			errs = append(errs, exportErrors...)
 		}
 
 		if dropped != 0 {
@@ -252,13 +252,37 @@ func (prwe *PrwExporter) handleSummaryMetric(tsMap map[string]*prompb.TimeSeries
 }
 
 // export sends a Snappy-compressed WriteRequest containing TimeSeries to a remote write endpoint in order
-func (prwe *PrwExporter) export(ctx context.Context, tsMap map[string]*prompb.TimeSeries) error {
-	// Calls the helper function to convert the TsMap to the desired format
-	req, err := wrapTimeSeries(tsMap)
+func (prwe *PrwExporter) export(ctx context.Context, tsMap map[string]*prompb.TimeSeries) []error {
+	var errs []error
+	// Calls the helper function to convert and batch the TsMap to the desired format
+	requests, err := wrapAndBatchTimeSeries(tsMap)
 	if err != nil {
-		return consumererror.Permanent(err)
+		errs = append(errs, err)
+		return errs
 	}
 
+	var wg sync.WaitGroup
+	var mtx sync.Mutex
+
+	// performs the batched requests concurrently
+	for _, request := range requests {
+		wg.Add(1)
+		go func(req *prompb.WriteRequest) {
+			requestError := prwe.execute(ctx, req)
+			if requestError != nil {
+				mtx.Lock()
+				errs = append(errs, consumererror.Permanent(requestError))
+				mtx.Unlock()
+			}
+			wg.Done()
+		}(request)
+	}
+	wg.Wait()
+
+	return errs
+}
+
+func (prwe *PrwExporter) execute(ctx context.Context, req *prompb.WriteRequest) error {
 	// Uses proto.Marshal to convert the WriteRequest into bytes array
 	data, err := proto.Marshal(req)
 	if err != nil {

exporter/prometheusremotewriteexporter/exporter_test.go

@@ -227,17 +227,19 @@ func Test_export(t *testing.T) {
 			if !tt.serverUp {
 				server.Close()
 			}
-			err := runExportPipeline(ts1, serverURL)
+			errs := runExportPipeline(ts1, serverURL)
 			if tt.returnError {
-				assert.Error(t, err)
+				assert.Error(t, errs[0])
 				return
 			}
-			assert.NoError(t, err)
+			assert.Len(t, errs, 0)
 		})
 	}
 }
 
-func runExportPipeline(ts *prompb.TimeSeries, endpoint *url.URL) error {
+func runExportPipeline(ts *prompb.TimeSeries, endpoint *url.URL) []error {
+	var errs []error
+
 	// First we will construct a TimeSeries array from the testutils package
 	testmap := make(map[string]*prompb.TimeSeries)
 	testmap["test"] = ts
@@ -246,10 +248,11 @@ func runExportPipeline(ts *prompb.TimeSeries, endpoint *url.URL) error {
 	// after this, instantiate a CortexExporter with the current HTTP client and endpoint set to passed in endpoint
 	prwe, err := NewPrwExporter("test", endpoint.String(), HTTPClient, map[string]string{})
 	if err != nil {
-		return err
+		errs = append(errs, err)
+		return errs
 	}
-	err = prwe.export(context.Background(), testmap)
-	return err
+	errs = append(errs, prwe.export(context.Background(), testmap)...)
+	return errs
 }
 
 // Test_PushMetrics checks the number of TimeSeries received by server and the number of metrics dropped is the same as

exporter/prometheusremotewriteexporter/helper.go

@@ -23,6 +23,7 @@ import (
 	"time"
 	"unicode"
 
+	"github.com/gogo/protobuf/proto"
 	"github.com/prometheus/prometheus/prompb"
 
 	"go.opentelemetry.io/collector/consumer/pdata"
@@ -41,6 +42,10 @@ const (
 	totalStr    = "total"
 	delimeter   = "_"
 	keyStr      = "key"
+
+	checkByteSizeInterval = 5000
+	maxBatchByteSize      = 3000000
+	maxBatchMetric        = 40000
 )
 
 // ByLabelName enables the usage of sort.Sort() with a slice of labels
@@ -215,21 +220,47 @@ func getPromMetricName(metric *otlp.Metric, ns string) string {
 	return sanitize(b.String())
 }
 
-// Simple helper function that takes the <Signature String - *TimeSeries> map
-// and creates a WriteRequest from the struct -- can move to the helper.go file
-func wrapTimeSeries(tsMap map[string]*prompb.TimeSeries) (*prompb.WriteRequest, error) {
+// Helper function that takes the <Signature String - *TimeSeries> map
+// and creates a batch of WriteRequest from the struct
+func wrapAndBatchTimeSeries(tsMap map[string]*prompb.TimeSeries) ([]*prompb.WriteRequest, error) {
 	if len(tsMap) == 0 {
 		return nil, errors.New("invalid tsMap: cannot be empty map")
 	}
+
+	var requests []*prompb.WriteRequest
 	var tsArray []prompb.TimeSeries
+	missingLastBatch := true
+
 	for _, v := range tsMap {
 		tsArray = append(tsArray, *v)
+		missingLastBatch = true
+
+		if len(tsArray)%checkByteSizeInterval == 0 {
+			wrapped := prompb.WriteRequest{
+				Timeseries: tsArray,
+				// Other parameters of the WriteRequest are unnecessary for our Export
+			}
+
+			if b := proto.Size(&wrapped); b >= maxBatchByteSize || len(tsArray) >= maxBatchMetric {
+				requests = append(requests, &wrapped)
+				tsArray = make([]prompb.TimeSeries, 0)
+
+				missingLastBatch = false
+			}
+		}
 	}
-	wrapped := prompb.WriteRequest{
-		Timeseries: tsArray,
-		// Other parameters of the WriteRequest are unnecessary for our Export
+
+	if missingLastBatch {
+		wrapped := prompb.WriteRequest{
+			Timeseries: tsArray,
+			// Other parameters of the WriteRequest are unnecessary for our Export
+		}
+
+		requests = append(requests, &wrapped)
+		tsArray = make([]prompb.TimeSeries, 0)
 	}
-	return &wrapped, nil
+
+	return requests, nil
 }
 
 // convertTimeStamp converts OTLP timestamp in ns to timestamp in ms