Concurrent flush for batches

codebien · olegbespalov · commit 739739af21a0 · 2023-07-20T13:30:13.000+02:00
diff --git a/output/cloud/expv2/flush.go b/output/cloud/expv2/flush.go
@@ -1,6 +1,8 @@
 package expv2
 
 import (
+	"math"
+	"sync"
 	"time"
 
 	"github.com/sirupsen/logrus"
@@ -21,6 +23,7 @@ type metricsFlusher struct {
 	discardedLabels            map[string]struct{}
 	aggregationPeriodInSeconds uint32
 	maxSeriesInBatch           int
+	batchPushConcurrency       int
 }
 
 // flush flushes the queued buckets sending them to the remote Cloud service.
@@ -43,45 +46,93 @@ func (f *metricsFlusher) flush() error {
 	// the metricSetBuilder is used for doing it during the traverse of the buckets.
 
 	var (
-		seriesCount  int
-		batchesCount int
-		start        = time.Now()
+		start       = time.Now()
+		batches     []metricSetBuilder
+		seriesCount int
 	)
 
 	defer func() {
 		f.logger.
 			WithField("t", time.Since(start)).
 			WithField("series", seriesCount).
 			WithField("buckets", len(buckets)).
-			WithField("batches", batchesCount).Debug("Flush the queued buckets")
+			WithField("batches", len(batches)).Debug("Flush the queued buckets")
 	}()
 
 	msb := newMetricSetBuilder(f.testRunID, f.aggregationPeriodInSeconds)
 	for i := 0; i < len(buckets); i++ {
 		for timeSeries, sink := range buckets[i].Sinks {
 			msb.addTimeSeries(buckets[i].Time, timeSeries, sink)
-			if len(msb.seriesIndex) < f.maxSeriesInBatch {
+			if len(msb.seriesIndex) <= f.maxSeriesInBatch {
 				continue
 			}
 
-			// we hit the batch size, let's flush
-			batchesCount++
+			// We hit the batch size, let's flush
 			seriesCount += len(msb.seriesIndex)
-			if err := f.push(msb); err != nil {
-				return err
-			}
+			batches = append(batches, msb)
+
+			// Reset the builder
 			msb = newMetricSetBuilder(f.testRunID, f.aggregationPeriodInSeconds)
 		}
 	}
 
-	if len(msb.seriesIndex) < 1 {
-		return nil
-	}
-
 	// send the last (or the unique) MetricSet chunk to the remote service
-	batchesCount++
 	seriesCount += len(msb.seriesIndex)
-	return f.push(msb)
+	batches = append(batches, msb)
+
+	return f.flushBatches(batches)
+}
+
+func (f *metricsFlusher) flushBatches(batches []metricSetBuilder) error {
+	var (
+		wg   = sync.WaitGroup{}
+		errs = make(chan error)
+		done = make(chan struct{})
+		stop = make(chan struct{})
+
+		workers   = int(math.Min(float64(len(batches)), float64(f.batchPushConcurrency)))
+		chunkSize = len(batches) / workers
+	)
+
+	wg.Add(workers)
+	for workersIndex := 0; workersIndex < workers; workersIndex++ {
+		offset := (workersIndex * chunkSize)
+		end := offset + chunkSize
+		if workersIndex == workers-1 {
+			end = len(batches)
+		}
+
+		go func(chunk []metricSetBuilder) {
+			defer wg.Done()
+
+			for i := 0; i < len(chunk); i++ {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				if err := f.push(chunk[i]); err != nil {
+					errs <- err
+				}
+			}
+		}(batches[offset:end])
+	}
+
+	go func() {
+		defer close(done)
+		wg.Wait()
+	}()
+
+	for {
+		select {
+		case err := <-errs:
+			close(stop)
+			<-done
+			return err
+		case <-done:
+			return nil
+		}
+	}
 }
 
 // push sends the metric set to the remote service.
diff --git a/output/cloud/expv2/flush_test.go b/output/cloud/expv2/flush_test.go
@@ -2,6 +2,8 @@ package expv2
 
 import (
 	"strconv"
+	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/sirupsen/logrus"
@@ -56,11 +58,12 @@ func TestMetricsFlusherFlushInBatchWithinBucket(t *testing.T) {
 		bq := &bucketQ{}
 		pm := &pusherMock{}
 		mf := metricsFlusher{
-			bq:               bq,
-			client:           pm,
-			logger:           logger,
-			discardedLabels:  make(map[string]struct{}),
-			maxSeriesInBatch: 3,
+			bq:                   bq,
+			client:               pm,
+			logger:               logger,
+			discardedLabels:      make(map[string]struct{}),
+			maxSeriesInBatch:     3,
+			batchPushConcurrency: 5,
 		}
 
 		bq.buckets = make([]timeBucket, 0, tc.series)
@@ -78,7 +81,7 @@ func TestMetricsFlusherFlushInBatchWithinBucket(t *testing.T) {
 		bq.Push([]timeBucket{{Time: 1, Sinks: sinks}})
 		err := mf.flush()
 		require.NoError(t, err)
-		assert.Equal(t, tc.expFlushCalls, pm.pushCalled)
+		assert.Equal(t, tc.expFlushCalls, pm.timesCalled())
 	}
 }
 
@@ -101,11 +104,12 @@ func TestMetricsFlusherFlushInBatchAcrossBuckets(t *testing.T) {
 		bq := &bucketQ{}
 		pm := &pusherMock{}
 		mf := metricsFlusher{
-			bq:               bq,
-			client:           pm,
-			logger:           logger,
-			discardedLabels:  make(map[string]struct{}),
-			maxSeriesInBatch: 3,
+			bq:                   bq,
+			client:               pm,
+			logger:               logger,
+			discardedLabels:      make(map[string]struct{}),
+			maxSeriesInBatch:     3,
+			batchPushConcurrency: 5,
 		}
 
 		bq.buckets = make([]timeBucket, 0, tc.series)
@@ -127,7 +131,7 @@ func TestMetricsFlusherFlushInBatchAcrossBuckets(t *testing.T) {
 
 		err := mf.flush()
 		require.NoError(t, err)
-		assert.Equal(t, tc.expFlushCalls, pm.pushCalled)
+		assert.Equal(t, tc.expFlushCalls, pm.timesCalled())
 	}
 }
 
@@ -136,21 +140,25 @@ func TestFlushWithReservedLabels(t *testing.T) {
 
 	logger, hook := testutils.NewLoggerWithHook(t)
 
+	mutex := sync.Mutex{}
 	collected := make([]*pbcloud.MetricSet, 0)
 
 	bq := &bucketQ{}
 	pm := &pusherMock{
 		hook: func(ms *pbcloud.MetricSet) {
+			mutex.Lock()
 			collected = append(collected, ms)
+			mutex.Unlock()
 		},
 	}
 
 	mf := metricsFlusher{
-		bq:               bq,
-		client:           pm,
-		maxSeriesInBatch: 2,
-		logger:           logger,
-		discardedLabels:  make(map[string]struct{}),
+		bq:                   bq,
+		client:               pm,
+		maxSeriesInBatch:     2,
+		logger:               logger,
+		discardedLabels:      make(map[string]struct{}),
+		batchPushConcurrency: 5,
 	}
 
 	r := metrics.NewRegistry()
@@ -186,7 +194,7 @@ func TestFlushWithReservedLabels(t *testing.T) {
 	require.NoError(t, err)
 
 	loglines := hook.Drain()
-	assert.Equal(t, 1, len(collected))
+	require.Len(t, collected, 1)
 
 	// check that warnings sown only once per label
 	assert.Len(t, testutils.FilterEntries(loglines, logrus.WarnLevel, "Tag __name__ has been discarded since it is reserved for Cloud operations."), 1)
@@ -207,14 +215,18 @@ func TestFlushWithReservedLabels(t *testing.T) {
 
 type pusherMock struct {
 	hook       func(*pbcloud.MetricSet)
-	pushCalled int
+	pushCalled int64
+}
+
+func (pm *pusherMock) timesCalled() int {
+	return int(atomic.LoadInt64(&pm.pushCalled))
 }
 
 func (pm *pusherMock) push(ms *pbcloud.MetricSet) error {
 	if pm.hook != nil {
 		pm.hook(ms)
 	}
 
-	pm.pushCalled++
+	atomic.AddInt64(&pm.pushCalled, 1)
 	return nil
 }
diff --git a/output/cloud/expv2/output.go b/output/cloud/expv2/output.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"runtime"
 	"strconv"
 	"sync"
 	"time"
@@ -112,6 +113,10 @@ func (o *Output) Start() error {
 		discardedLabels:            make(map[string]struct{}),
 		aggregationPeriodInSeconds: uint32(o.config.AggregationPeriod.TimeDuration().Seconds()),
 		maxSeriesInBatch:           int(o.config.MaxTimeSeriesInBatch.Int64),
+		// TODO: when the migration from v1 is over
+		// change the default of cloudapi.MetricPushConcurrency to use GOMAXPROCS(0)
+		// batchPushConcurrency: int(o.config.MetricPushConcurrency.Int64),
+		batchPushConcurrency: runtime.GOMAXPROCS(0),
 	}
 
 	o.runFlushWorkers()
@@ -178,12 +183,13 @@ func (o *Output) StopWithTestError(_ error) error {
 }
 
 func (o *Output) runFlushWorkers() {
+	t := time.NewTicker(o.config.MetricPushInterval.TimeDuration())
+
+	// TODO: drop it when we are sure of the new proposed architecture
 	// workers := o.config.MetricPushConcurrency.Int64
 	// Details: https://github.com/grafana/k6/issues/3192
 	workers := 1
 
-	t := time.NewTicker(o.config.MetricPushInterval.TimeDuration())
-
 	for i := 0; i < workers; i++ {
 		o.wg.Add(1)
 		go func() {
