2021w15 scale benchtool (#167)

* feat(benchtool): use buffer pool for remote-write runner

Signed-off-by: Jacob Lisi <[email protected]>

* fixup series chan

Signed-off-by: Jacob Lisi <[email protected]>

Author: jtlisi

pkg/bench/workload.go

@@ -2,14 +2,17 @@ package bench
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"hash/adler32"
 	"math/rand"
 	"strings"
+	"sync"
 	"text/template"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/prometheus/prompb"
 )
 
@@ -70,7 +73,11 @@ type writeWorkload struct {
 	totalSeries        int
 	totalSeriesTypeMap map[SeriesType]int
 
+	missedIterations prometheus.Counter
+
 	options WriteDesc
+
+	seriesBufferChan chan []prompb.TimeSeries
 }
 
 func newWriteWorkload(workloadDesc WorkloadDesc, reg prometheus.Registerer) *writeWorkload {
@@ -132,6 +139,14 @@ func newWriteWorkload(workloadDesc WorkloadDesc, reg prometheus.Registerer) *wri
 		totalSeries:        totalSeries,
 		totalSeriesTypeMap: totalSeriesTypeMap,
 		options:            workloadDesc.Write,
+
+		missedIterations: promauto.With(reg).NewCounter(
+			prometheus.CounterOpts{
+				Namespace: "benchtool",
+				Name:      "write_iterations_late_total",
+				Help:      "Number of write intervals started late because the previous interval did not complete in time.",
+			},
+		),
 	}
 }
 
@@ -194,6 +209,91 @@ func (w *writeWorkload) generateTimeSeries(id string, t time.Time) []prompb.Time
 	return timeseries
 }
 
+type batchReq struct {
+	batch   []prompb.TimeSeries
+	wg      *sync.WaitGroup
+	putBack chan []prompb.TimeSeries
+}
+
+func (w *writeWorkload) getSeriesBuffer(ctx context.Context) []prompb.TimeSeries {
+	select {
+	case <-ctx.Done():
+		return nil
+	case seriesBuffer := <-w.seriesBufferChan:
+		return seriesBuffer[:0]
+	}
+}
+
+func (w *writeWorkload) generateWriteBatch(ctx context.Context, id string, numBuffers int, seriesChan chan batchReq) error {
+	w.seriesBufferChan = make(chan []prompb.TimeSeries, numBuffers)
+	for i := 0; i < numBuffers; i++ {
+		w.seriesBufferChan <- make([]prompb.TimeSeries, 0, w.options.BatchSize)
+	}
+
+	seriesBuffer := w.getSeriesBuffer(ctx)
+	ticker := time.NewTicker(w.options.Interval)
+
+	defer close(seriesChan)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		case timeNow := <-ticker.C:
+			now := timeNow.UnixNano() / int64(time.Millisecond)
+			wg := &sync.WaitGroup{}
+			for replicaNum := 0; replicaNum < w.replicas; replicaNum++ {
+				replicaLabel := prompb.Label{Name: "bench_replica", Value: fmt.Sprintf("replica-%05d", replicaNum)}
+				idLabel := prompb.Label{Name: "bench_id", Value: id}
+				for _, series := range w.series {
+					var value float64
+					switch series.seriesType {
+					case GaugeZero:
+						value = 0
+					case GaugeRandom:
+						value = rand.Float64()
+					case CounterOne:
+						value = series.lastValue + 1
+					case CounterRandom:
+						value = series.lastValue + float64(rand.Int())
+					default:
+						return fmt.Errorf("unknown series type %v", series.seriesType)
+					}
+					series.lastValue = value
+					for _, labelSet := range series.labelSets {
+						if len(seriesBuffer) == w.options.BatchSize {
+							wg.Add(1)
+							seriesChan <- batchReq{seriesBuffer, wg, w.seriesBufferChan}
+							seriesBuffer = w.getSeriesBuffer(ctx)
+						}
+						newLabelSet := make([]prompb.Label, len(labelSet)+2)
+						copy(newLabelSet, labelSet)
+
+						newLabelSet[len(newLabelSet)-2] = replicaLabel
+						newLabelSet[len(newLabelSet)-1] = idLabel
+						seriesBuffer = append(seriesBuffer, prompb.TimeSeries{
+							Labels: newLabelSet,
+							Samples: []prompb.Sample{{
+								Timestamp: now,
+								Value:     value,
+							}},
+						})
+					}
+				}
+			}
+			if len(seriesBuffer) > 0 {
+				wg.Add(1)
+				seriesChan <- batchReq{seriesBuffer, wg, w.seriesBufferChan}
+				seriesBuffer = w.getSeriesBuffer(ctx)
+			}
+			wg.Wait()
+			if time.Since(timeNow) > w.options.Interval {
+				w.missedIterations.Inc()
+			}
+		}
+	}
+}
+
 type queryWorkload struct {
 	queries []query
 }

pkg/bench/write_runner.go

@@ -55,8 +55,7 @@ type WriteBenchmarkRunner struct {
 	reg    prometheus.Registerer
 	logger log.Logger
 
-	requestDuration  *prometheus.HistogramVec
-	missedIterations prometheus.Counter
+	requestDuration *prometheus.HistogramVec
 }
 
 func NewWriteBenchmarkRunner(id string, tenantName string, cfg WriteBenchConfig, workload *writeWorkload, logger log.Logger, reg prometheus.Registerer) (*WriteBenchmarkRunner, error) {
@@ -82,13 +81,6 @@ func NewWriteBenchmarkRunner(id string, tenantName string, cfg WriteBenchConfig,
 			},
 			[]string{"code"},
 		),
-		missedIterations: promauto.With(reg).NewCounter(
-			prometheus.CounterOpts{
-				Namespace: "benchtool",
-				Name:      "write_iterations_late_total",
-				Help:      "Number of write intervals started late because the previous interval did not complete in time.",
-			},
-		),
 	}
 
 	// Resolve an initial set of distributor addresses
@@ -143,50 +135,16 @@ func (w *WriteBenchmarkRunner) Run(ctx context.Context) error {
 	// Start a loop to re-resolve addresses every 5 minutes
 	go w.resolveAddrsLoop(ctx)
 
+	// Run replicas * 10 write client workers.
+	// This number will also be used for the number of series buffers to store at once.
+	numWorkers := w.workload.replicas * 10
+
 	batchChan := make(chan batchReq, 10)
-	for i := 0; i < w.workload.replicas*10; i++ {
+	for i := 0; i < numWorkers; i++ {
 		go w.writeWorker(batchChan)
 	}
 
-	ticker := time.NewTicker(w.workload.options.Interval)
-	for {
-		select {
-		case <-ctx.Done():
-			close(batchChan)
-			return nil
-		case now := <-ticker.C:
-			timeseries := w.workload.generateTimeSeries(w.id, now)
-			batchSize := w.workload.options.BatchSize
-			var batches [][]prompb.TimeSeries
-			if batchSize < len(timeseries) {
-				batches = make([][]prompb.TimeSeries, 0, (len(timeseries)+batchSize-1)/batchSize)
-
-				level.Info(w.logger).Log("msg", "sending timeseries", "num_series", strconv.Itoa(len(timeseries)))
-				for batchSize < len(timeseries) {
-					timeseries, batches = timeseries[batchSize:], append(batches, timeseries[0:batchSize:batchSize])
-				}
-			} else {
-				batches = [][]prompb.TimeSeries{timeseries}
-			}
-
-			wg := &sync.WaitGroup{}
-			for _, batch := range batches {
-				reqBatch := batch
-				wg.Add(1)
-				batchChan <- batchReq{reqBatch, wg}
-			}
-
-			wg.Wait()
-			if time.Since(now) > w.workload.options.Interval {
-				w.missedIterations.Inc()
-			}
-		}
-	}
-}
-
-type batchReq struct {
-	batch []prompb.TimeSeries
-	wg    *sync.WaitGroup
+	return w.workload.generateWriteBatch(ctx, w.id, numWorkers+10, batchChan)
 }
 
 func (w *WriteBenchmarkRunner) writeWorker(batchChan chan batchReq) {
@@ -195,6 +153,9 @@ func (w *WriteBenchmarkRunner) writeWorker(batchChan chan batchReq) {
 		if err != nil {
 			level.Warn(w.logger).Log("msg", "unable to send batch", "err", err)
 		}
+
+		// put back the series buffer
+		batchReq.putBack <- batchReq.batch
 		batchReq.wg.Done()
 	}
 }