Add concurrency to block-builder wal conversion and flushing

Author: mdisibio

modules/blockbuilder/blockbuilder_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/go-kit/log"
 	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/ring"
 	"github.com/grafana/dskit/services"
@@ -328,7 +329,12 @@ type ownEverythingSharder struct{}
 func (o *ownEverythingSharder) Owns(string) bool { return true }
 
 func newStore(ctx context.Context, t testing.TB) storage.Store {
+	return newStoreWithLogger(ctx, t, test.NewTestingLogger(t))
+}
+
+func newStoreWithLogger(ctx context.Context, t testing.TB, log log.Logger) storage.Store {
 	tmpDir := t.TempDir()
+
 	s, err := storage.NewStore(storage.Config{
 		Trace: tempodb.Config{
 			Backend: backend.Local,
@@ -348,7 +354,7 @@ func newStore(ctx context.Context, t testing.TB) storage.Store {
 			},
 			BlocklistPoll: 5 * time.Second,
 		},
-	}, nil, test.NewTestingLogger(t))
+	}, nil, log)
 	require.NoError(t, err)
 
 	s.EnablePolling(ctx, &ownEverythingSharder{})
@@ -483,14 +489,22 @@ func requireLastCommitEquals(t testing.TB, ctx context.Context, client *kgo.Clie
 func BenchmarkBlockBuilder(b *testing.B) {
 	var (
 		ctx        = context.Background()
+		logger     = log.NewNopLogger()
 		_, address = testkafka.CreateCluster(b, 1, testTopic)
-		store      = newStore(ctx, b)
+		store      = newStoreWithLogger(ctx, b, logger)
 		cfg        = blockbuilderConfig(b, address)
 		client     = newKafkaClient(b, cfg.IngestStorageConfig.Kafka)
 	)
 
 	cfg.ConsumeCycleDuration = 1 * time.Hour
 
+	bb := New(cfg, logger, newPartitionRingReader(), &mockOverrides{}, store)
+	defer bb.stopping(nil)
+
+	// Startup (without starting the background consume cycle)
+	err := bb.starting(ctx)
+	require.NoError(b, err)
+
 	b.ResetTimer()
 
 	for i := 0; i < b.N; i++ {
@@ -508,12 +522,6 @@ func BenchmarkBlockBuilder(b *testing.B) {
 
 		b.ResetTimer()
 
-		bb := New(cfg, test.NewTestingLogger(b), newPartitionRingReader(), &mockOverrides{}, store)
-
-		// Startup (without starting the background consume cycle)
-		err := bb.starting(ctx)
-		require.NoError(b, err)
-
 		err = bb.consume(ctx)
 		require.NoError(b, err)
 

modules/blockbuilder/tenant_store.go

@@ -12,6 +12,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/grafana/tempo/modules/blockbuilder/util"
 	"github.com/grafana/tempo/modules/overrides"
+	"github.com/grafana/tempo/pkg/boundedwaitgroup"
 	"github.com/grafana/tempo/pkg/livetraces"
 	"github.com/grafana/tempo/pkg/model"
 	"github.com/grafana/tempo/pkg/tempopb"
@@ -23,6 +24,7 @@ import (
 	"github.com/grafana/tempo/tempodb/wal"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"go.uber.org/atomic"
 )
 
 var metricBlockBuilderFlushedBlocks = promauto.NewCounterVec(
@@ -33,7 +35,10 @@ var metricBlockBuilderFlushedBlocks = promauto.NewCounterVec(
 	}, []string{"tenant"},
 )
 
-const reasonTraceTooLarge = "trace_too_large"
+const (
+	reasonTraceTooLarge = "trace_too_large"
+	flushConcurrency    = 4
+)
 
 // TODO - This needs locking
 type tenantStore struct {
@@ -192,34 +197,62 @@ func (s *tenantStore) Flush(ctx context.Context, store tempodb.Writer) error {
 	s.blocksMtx.Lock()
 	defer s.blocksMtx.Unlock()
 
-	completeBlocks := make([]tempodb.WriteableBlock, 0, len(s.walBlocks))
-	// Write all blocks
-	for _, block := range s.walBlocks {
-		completeBlock, err := s.buildWriteableBlock(ctx, block)
-		if err != nil {
-			return err
-		}
+	var (
+		completeBlocks = make([]tempodb.WriteableBlock, len(s.walBlocks))
+		jobErr         = atomic.NewError(nil)
+		wg             = boundedwaitgroup.New(flushConcurrency)
+	)
+
+	// Convert WALs to backend blocks
+	for i, block := range s.walBlocks {
+		wg.Add(1)
+		go func(i int, block common.WALBlock) {
+			defer wg.Done()
+
+			completeBlock, err := s.buildWriteableBlock(ctx, block)
+			if err != nil {
+				jobErr.Store(err)
+				return
+			}
 
-		err = block.Clear()
-		if err != nil {
-			return err
-		}
+			err = block.Clear()
+			if err != nil {
+				jobErr.Store(err)
+				return
+			}
 
-		completeBlocks = append(completeBlocks, completeBlock)
+			completeBlocks[i] = completeBlock
+		}(i, block)
+	}
+
+	wg.Wait()
+	if err := jobErr.Load(); err != nil {
+		return err
 	}
 
-	level.Info(s.logger).Log("msg", "writing blocks to storage", "num_blocks", len(completeBlocks))
 	// Write all blocks to the store
+	level.Info(s.logger).Log("msg", "writing blocks to storage", "num_blocks", len(completeBlocks))
 	for _, block := range completeBlocks {
-		level.Info(s.logger).Log("msg", "writing block to storage", "block_id", block.BlockMeta().BlockID.String())
-		if err := store.WriteBlock(ctx, block); err != nil {
-			return err
-		}
-		metricBlockBuilderFlushedBlocks.WithLabelValues(s.tenantID).Inc()
+		wg.Add(1)
+		go func(block tempodb.WriteableBlock) {
+			defer wg.Done()
+			level.Info(s.logger).Log("msg", "writing block to storage", "block_id", block.BlockMeta().BlockID.String())
+			if err := store.WriteBlock(ctx, block); err != nil {
+				jobErr.Store(err)
+				return
+			}
 
-		if err := s.wal.LocalBackend().ClearBlock((uuid.UUID)(block.BlockMeta().BlockID), s.tenantID); err != nil {
-			return err
-		}
+			metricBlockBuilderFlushedBlocks.WithLabelValues(s.tenantID).Inc()
+
+			if err := s.wal.LocalBackend().ClearBlock((uuid.UUID)(block.BlockMeta().BlockID), s.tenantID); err != nil {
+				jobErr.Store(err)
+			}
+		}(block)
+	}
+
+	wg.Wait()
+	if err := jobErr.Load(); err != nil {
+		return err
 	}
 
 	// Clear the blocks