Block-builder: adopt concurrent fetcher from ingest storage. (#12222)

Reuse ingest-storage's ConcurrentFetchers inside the block-builder to
consume faster from Kafka when `kafka.fetch-concurrency-max` is given in
the flags.

Author: seizethedave

CHANGELOG.md

@@ -52,6 +52,7 @@
 * [ENHANCEMENT] Querier: Include more information about inflight queries in the activity tracker. A querier logs this information after it restarts following a crash. #12526
 * [ENHANCEMENT] Ingester: Add experimental `-blocks-storage.tsdb.index-lookup-planning-comparison-portion` flag to enable mirrored chunk querier comparison between queries with and without index lookup planning. #12460
 * [ENHANCEMENT] Ruler: Add native histogram version of `cortex_ruler_sync_rules_duration_seconds`. #12628
+* [ENHANCEMENT] Block-builder: Implement concurrent consumption within a job when `-ingest-storage.kafka.fetch-concurrency-max` is given. #12222
 * [ENHANCEMENT] Query-frontend: Labels query optimizer is no longer experimental and is enabled by default. It can be disabled with `-query-frontend.labels-query-optimizer-enabled=false` CLI flag. #12606
 * [ENHANCEMENT] Distributor: Add value length to "label value too long" error. #12583
 * [ENHANCEMENT] Distributor: The metric `cortex_distributor_uncompressed_request_body_size_bytes` now differentiates by the handler serving the request. #12661

pkg/blockbuilder/blockbuilder.go

@@ -20,9 +20,9 @@ import (
 	"github.com/prometheus/prometheus/tsdb"
 	"github.com/thanos-io/objstore"
 	"github.com/twmb/franz-go/pkg/kgo"
+	"github.com/twmb/franz-go/plugin/kprom"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"go.opentelemetry.io/otel"
-	"go.uber.org/atomic"
 	"google.golang.org/grpc"
 
 	"github.com/grafana/mimir/pkg/blockbuilder/schedulerpb"
@@ -49,11 +49,11 @@ type BlockBuilder struct {
 	schedulerClient schedulerpb.SchedulerClient
 	schedulerConn   *grpc.ClientConn
 
-	// the current job iteration number. For tests.
-	jobIteration atomic.Int64
-
 	blockBuilderMetrics   blockBuilderMetrics
 	tsdbBuilderMetrics    tsdbBuilderMetrics
+	readerMetrics         *ingest.ReaderMetrics
+	readerMetricsSource   swappableReaderMetricsSource
+	kpromMetrics          *kprom.Metrics
 	pusherConsumerMetrics *ingest.PusherConsumerMetrics
 }
 
@@ -74,13 +74,25 @@ func newWithSchedulerClient(
 	limits *validation.Overrides,
 	schedulerClient schedulerpb.SchedulerClient,
 ) (*BlockBuilder, error) {
+	kpm := ingest.NewKafkaReaderClientMetrics(ingest.ReaderMetricsPrefix, "block-builder", reg)
+	readerMetricsSource := swappableReaderMetricsSource{&zeroReaderMetricsSource{}}
+
+	var readerMetrics *ingest.ReaderMetrics
+	if cfg.Kafka.FetchConcurrencyMax > 0 {
+		m := ingest.NewReaderMetrics(reg, readerMetricsSource, cfg.Kafka.Topic, kpm)
+		readerMetrics = &m
+	}
+
 	b := &BlockBuilder{
 		cfg:                   cfg,
 		logger:                logger,
 		register:              reg,
 		limits:                limits,
 		blockBuilderMetrics:   newBlockBuilderMetrics(reg),
 		tsdbBuilderMetrics:    newTSDBBBuilderMetrics(reg),
+		readerMetrics:         readerMetrics,
+		readerMetricsSource:   readerMetricsSource,
+		kpromMetrics:          kpm,
 		pusherConsumerMetrics: ingest.NewPusherConsumerMetrics(reg),
 	}
 
@@ -100,7 +112,6 @@ func newWithSchedulerClient(
 	}
 
 	b.Service = services.NewBasicService(b.starting, b.running, b.stopping)
-
 	return b, nil
 }
 
@@ -146,7 +157,7 @@ func (b *BlockBuilder) starting(context.Context) (err error) {
 
 	b.kafkaClient, err = ingest.NewKafkaReaderClient(
 		b.cfg.Kafka,
-		ingest.NewKafkaReaderClientMetrics(ingest.ReaderMetricsPrefix, "block-builder", b.register),
+		b.kpromMetrics,
 		b.logger,
 	)
 	if err != nil {
@@ -168,7 +179,6 @@ func (b *BlockBuilder) stopping(_ error) error {
 
 // running learns about the jobs from a block-builder-scheduler, and consumes one job at a time.
 func (b *BlockBuilder) running(ctx context.Context) error {
-
 	// Block-builder attempts to complete the current job when a shutdown
 	// request is received.
 	// To enable this, we create a child context whose cancellation signal is
@@ -200,7 +210,7 @@ func (b *BlockBuilder) running(ctx context.Context) error {
 		}
 
 		// Once we've gotten a job, we attempt to complete it even if the context is cancelled.
-		if _, err := b.consumeJob(graceCtx, key, spec); err != nil {
+		if err := b.consumeJob(graceCtx, key, spec); err != nil {
 			level.Error(b.logger).Log("msg", "failed to consume job", "job_id", key.Id, "epoch", key.Epoch, "err", err)
 
 			if err := b.schedulerClient.FailJob(key); err != nil {
@@ -214,13 +224,11 @@ func (b *BlockBuilder) running(ctx context.Context) error {
 		if err := b.schedulerClient.CompleteJob(key); err != nil {
 			level.Error(b.logger).Log("msg", "failed to complete job", "job_id", key.Id, "epoch", key.Epoch, "err", err)
 		}
-
-		b.jobIteration.Inc()
 	}
 }
 
 // consumeJob performs block consumption from Kafka into object storage based on the given job spec.
-func (b *BlockBuilder) consumeJob(ctx context.Context, key schedulerpb.JobKey, spec schedulerpb.JobSpec) (lastOffset int64, err error) {
+func (b *BlockBuilder) consumeJob(ctx context.Context, key schedulerpb.JobKey, spec schedulerpb.JobSpec) (err error) {
 	defer func(start time.Time) {
 		success := "true"
 		if err != nil {
@@ -243,6 +251,84 @@ func (b *BlockBuilder) consumeJob(ctx context.Context, key schedulerpb.JobKey, s
 	return b.consumePartitionSection(ctx, logger, consumer, builder, spec.Partition, spec.StartOffset, spec.EndOffset)
 }
 
+type fetchPoller interface {
+	PollFetches(context.Context) kgo.Fetches
+}
+
+type fetchWrapper struct {
+	fetchers *ingest.ConcurrentFetchers
+}
+
+func (f *fetchWrapper) PollFetches(ctx context.Context) kgo.Fetches {
+	fetch, _ := f.fetchers.PollFetches(ctx)
+	return fetch
+}
+
+var _ fetchPoller = (*fetchWrapper)(nil)
+
+// swappableReaderMetricsSource is a ReaderMetricsSource that can be swapped out at runtime.
+type swappableReaderMetricsSource struct {
+	ingest.ReaderMetricsSource
+}
+
+func (s *swappableReaderMetricsSource) set(metricsSource ingest.ReaderMetricsSource) {
+	s.ReaderMetricsSource = metricsSource
+}
+
+type zeroReaderMetricsSource struct{}
+
+func (z *zeroReaderMetricsSource) BufferedBytes() int64           { return 0 }
+func (z *zeroReaderMetricsSource) BufferedRecords() int64         { return 0 }
+func (z *zeroReaderMetricsSource) EstimatedBytesPerRecord() int64 { return 0 }
+
+var _ ingest.ReaderMetricsSource = (*zeroReaderMetricsSource)(nil)
+
+// newFetchers creates a new concurrent fetcher, retrying until it succeeds or the context is cancelled.
+// The returned error is the last error encountered.
+func (b *BlockBuilder) newFetchers(ctx context.Context, logger log.Logger, partition int32, startOffset int64) (*ingest.ConcurrentFetchers, error) {
+	if b.readerMetrics == nil {
+		panic("readerMetrics should be non-nil when concurrent fetchers are used")
+	}
+
+	boff := backoff.New(ctx, backoff.Config{
+		MinBackoff: 100 * time.Millisecond,
+		MaxBackoff: 5 * time.Second,
+		MaxRetries: 10,
+	})
+
+	var lastError error
+
+	for boff.Ongoing() {
+		f, ferr := ingest.NewConcurrentFetchers(
+			ctx,
+			b.kafkaClient,
+			logger,
+			b.cfg.Kafka.Topic,
+			partition,
+			startOffset,
+			b.cfg.Kafka.FetchConcurrencyMax,
+			int32(b.cfg.Kafka.MaxBufferedBytes),
+			b.cfg.Kafka.UseCompressedBytesAsFetchMaxBytes,
+			b.cfg.Kafka.FetchMaxWait,
+			nil, // Don't need a reader since we've provided the start offset.
+			ingest.OnRangeErrorAbort,
+			nil, // We're aborting on range error, so we don't need an offset reader.
+			backoff.Config{
+				MinBackoff: 100 * time.Millisecond,
+				MaxBackoff: 1 * time.Second,
+			},
+			b.readerMetrics)
+		if ferr == nil {
+			return f, nil
+		}
+		level.Warn(b.logger).Log("msg", "failed to create concurrent fetcher, probably retrying...", "err", ferr)
+		lastError = ferr
+		boff.Wait()
+	}
+
+	return nil, lastError
+}
+
 // consumePartitionSection is for the use of scheduler-based architecture.
 // startOffset is inclusive, endOffset is exclusive, and must be valid offsets and not something in the future (endOffset can be technically 1 offset in the future).
 // All the records and samples between these offsets will be consumed and put into a block.
@@ -254,8 +340,8 @@ func (b *BlockBuilder) consumePartitionSection(
 	builder *TSDBBuilder,
 	partition int32,
 	startOffset, endOffset int64,
-) (lastConsumedOffset int64, retErr error) {
-	lastConsumedOffset = startOffset
+) (retErr error) {
+	lastConsumedOffset := startOffset
 	if startOffset >= endOffset {
 		level.Info(logger).Log("msg", "nothing to consume")
 		return
@@ -287,29 +373,49 @@ func (b *BlockBuilder) consumePartitionSection(
 	})
 	defer b.kafkaClient.RemoveConsumePartitions(map[string][]int32{b.cfg.Kafka.Topic: {partition}})
 
+	var fetchPoller fetchPoller = b.kafkaClient
+
+	if b.cfg.Kafka.FetchConcurrencyMax > 0 {
+		f, ferr := b.newFetchers(ctx, logger, partition, startOffset)
+		if ferr != nil {
+			return fmt.Errorf("creating concurrent fetcher: %w", ferr)
+		}
+
+		b.readerMetricsSource.set(f)
+
+		f.Start(ctx)
+		defer f.Stop()
+
+		fetchPoller = &fetchWrapper{f}
+	}
+
 	level.Info(logger).Log("msg", "start consuming", "partition", partition, "start_offset", startOffset, "end_offset", endOffset)
 
-	var (
-		firstRecOffset = int64(-1)
-		lastRecOffset  = int64(-1)
-	)
+	firstRecOffset := int64(-1)
 
-	for lastRecOffset < endOffset-1 {
+	for lastConsumedOffset < endOffset-1 {
 		if err := context.Cause(ctx); err != nil {
-			return 0, err
+			return err
 		}
 
 		// PollFetches can return a non-failed fetch with zero records. In such a case, with only the fetches at hands,
 		// we cannot tell if the consumer has already reached the latest end of the partition, i.e. no more records to consume,
 		// or there is more data in the backlog, and we must retry the poll. That's why the consumer loop above has to guard
 		// the iterations against the endOffset, so it retries the polling up until the expected end of the partition is reached.
-		fetches := b.kafkaClient.PollFetches(ctx)
+		fetches := fetchPoller.PollFetches(ctx)
+		var fetchErr error
 		fetches.EachError(func(_ string, _ int32, err error) {
 			if !errors.Is(err, context.Canceled) {
 				level.Error(logger).Log("msg", "failed to fetch records", "err", err)
 				b.blockBuilderMetrics.fetchErrors.WithLabelValues(fmt.Sprintf("%d", partition)).Inc()
+				if fetchErr == nil {
+					fetchErr = err
+				}
 			}
 		})
+		if fetchErr != nil {
+			return fmt.Errorf("poll fetches: %w", fetchErr)
+		}
 
 		recordsAll := func(fetches kgo.Fetches) iter.Seq[*kgo.Record] {
 			return func(yield func(*kgo.Record) bool) {
@@ -328,15 +434,14 @@ func (b *BlockBuilder) consumePartitionSection(
 
 		records := recordsAll(fetches)
 		for rec := range records {
-			lastRecOffset = rec.Offset
+			lastConsumedOffset = rec.Offset
 			if firstRecOffset == -1 {
-				firstRecOffset = lastRecOffset
+				firstRecOffset = lastConsumedOffset
 			}
 		}
 
-		err := consumer.Consume(ctx, records)
-		if err != nil {
-			return 0, fmt.Errorf("consume records in partition %d: %w", partition, err)
+		if err := consumer.Consume(ctx, records); err != nil {
+			return fmt.Errorf("consume records in partition %d: %w", partition, err)
 		}
 	}
 
@@ -348,12 +453,12 @@ func (b *BlockBuilder) consumePartitionSection(
 	var err error
 	blockMetas, err = builder.CompactAndUpload(ctx, b.uploadBlocks)
 	if err != nil {
-		return 0, err
+		return err
 	}
 
 	// TODO: figure out a way to track the blockCounts metrics.
 
-	return lastRecOffset, nil
+	return nil
 }
 
 func (b *BlockBuilder) uploadBlocks(ctx context.Context, tenantID, dbDir string, metas []tsdb.BlockMeta) error {

pkg/blockbuilder/blockbuilder_test.go

@@ -130,7 +130,6 @@ func TestBlockBuilder(t *testing.T) {
 					labels.MustNewMatcher(labels.MatchRegexp, "foo", ".*"),
 				)
 			}
-
 		})
 	}
 }