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