wip

Author: grobinson-grafana

pkg/dataobj/consumer/config.go

@@ -4,21 +4,33 @@ import (
 	"flag"
 	"time"
 
+	"github.com/grafana/dskit/ring"
 	"github.com/grafana/loki/v3/pkg/dataobj/consumer/logsobj"
 	"github.com/grafana/loki/v3/pkg/dataobj/uploader"
+	"github.com/grafana/loki/v3/pkg/kafka/partitionring"
+	util_log "github.com/grafana/loki/v3/pkg/util/log"
 )
 
 type Config struct {
-	logsobj.BuilderConfig
-	UploaderConfig   uploader.Config `yaml:"uploader"`
-	IdleFlushTimeout time.Duration   `yaml:"idle_flush_timeout"`
+	BuilderConfig       logsobj.BuilderConfig `yaml:"builder,omitempty`
+	LifecyclerConfig    ring.LifecyclerConfig `yaml:"lifecycler,omitempty"`
+	PartitionRingConfig partitionring.Config  `yaml:"partition_ring"`
+	UploaderConfig      uploader.Config       `yaml:"uploader"`
+	IdleFlushTimeout    time.Duration         `yaml:"idle_flush_timeout"`
 }
 
 func (cfg *Config) Validate() error {
+	if err := cfg.BuilderConfig.Validate(); err != nil {
+		return err
+	}
+	if err := cfg.LifecyclerConfig.Validate(); err != nil {
+		return err
+	}
+	// The PartitionRingConfig does not have a validate method.
 	if err := cfg.UploaderConfig.Validate(); err != nil {
 		return err
 	}
-	return cfg.BuilderConfig.Validate()
+	return nil
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
@@ -27,6 +39,8 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 
 func (cfg *Config) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
 	cfg.BuilderConfig.RegisterFlagsWithPrefix(prefix, f)
+	cfg.LifecyclerConfig.RegisterFlagsWithPrefix(prefix, f, util_log.Logger)
+	cfg.PartitionRingConfig.RegisterFlagsWithPrefix(prefix, f)
 	cfg.UploaderConfig.RegisterFlagsWithPrefix(prefix, f)
 
 	f.DurationVar(&cfg.IdleFlushTimeout, prefix+"idle-flush-timeout", 60*60*time.Second, "The maximum amount of time to wait in seconds before flushing an object that is no longer receiving new writes")

pkg/dataobj/consumer/consumer.go

@@ -2,79 +2,48 @@ package consumer
 
 import (
 	"context"
-	"errors"
 	"sync"
-	"time"
 
 	"github.com/go-kit/log"
-	"github.com/go-kit/log/level"
-	"github.com/grafana/dskit/backoff"
-	"github.com/twmb/franz-go/pkg/kgo"
+	"github.com/grafana/loki/v3/pkg/kafka/partition"
 )
 
-// kafkaConsumer allows mocking of certain [kgo.Client] methods in tests.
-type kafkaConsumer interface {
-	PollFetches(context.Context) kgo.Fetches
-}
-
 // processor allows mocking of [partitionProcessor] in tests.
 type processor interface {
-	Append(records []*kgo.Record) bool
+	Append(records []partition.Record) bool
 }
 
 // consumer polls records from the Kafka topic and passes each record to
 // its indended processor.
 type consumer struct {
-	client    kafkaConsumer
 	logger    log.Logger
 	processor processor
 	mtx       sync.RWMutex
 }
 
 // newConsumer returns a new consumer.
-func newConsumer(client kafkaConsumer, processor processor, logger log.Logger) *consumer {
+func newConsumer(processor processor, logger log.Logger) *consumer {
 	return &consumer{
-		client:    client,
 		logger:    logger,
 		processor: processor,
 	}
 }
 
-// run starts the poll loop. It is stopped when either the context is canceled
-// or the kafka client is closed.
-func (c *consumer) Run(ctx context.Context) error {
-	b := backoff.New(ctx, backoff.Config{
-		MinBackoff: 100 * time.Millisecond,
-		MaxBackoff: time.Second,
-		MaxRetries: 0,
-	})
-	for b.Ongoing() {
-		select {
-		case <-ctx.Done():
-			return nil
-		default:
-			if err := c.pollFetches(ctx); err != nil {
-				if errors.Is(err, kgo.ErrClientClosed) {
-					return nil
-				}
-				level.Error(c.logger).Log("msg", "failed to poll fetches", "err", err.Error())
-				b.Wait()
-			}
+func (c *consumer) Start(ctx context.Context, recordsChan <-chan []partition.Record) func() {
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for records := range recordsChan {
+			c.processor.Append(records)
 		}
-	}
-	return nil
+	}()
+	return wg.Wait
 }
 
-func (c *consumer) pollFetches(ctx context.Context) error {
-	fetches := c.client.PollFetches(ctx)
-	// If the client is closed, or the context was canceled, return the error
-	// as no fetches were polled. We use this instead of [kgo.IsClientClosed]
-	// so we can also check if the context was canceled.
-	if err := fetches.Err0(); err != nil {
-		if errors.Is(err, kgo.ErrClientClosed) || errors.Is(err, context.Canceled) {
-			return err
-		}
+// newConsumerFactory returns a consumer factory.
+func newConsumerFactory(partitionProcessorFactory *partitionProcessorFactory) partition.ConsumerFactory {
+	return func(committer partition.Committer, logger log.Logger) (partition.Consumer, error) {
+		return newConsumer(partitionProcessorFactory.New(committer, logger), logger), nil
 	}
-	c.processor.Append(fetches.Records())
-	return nil
 }

pkg/dataobj/consumer/partition_processor.go

@@ -5,10 +5,8 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"strconv"
 	"sync"
 	"time"
-	"unicode/utf8"
 
 	"github.com/coder/quartz"
 	"github.com/go-kit/log"
@@ -24,6 +22,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/dataobj/metastore/multitenancy"
 	"github.com/grafana/loki/v3/pkg/dataobj/uploader"
 	"github.com/grafana/loki/v3/pkg/kafka"
+	"github.com/grafana/loki/v3/pkg/kafka/partition"
 	"github.com/grafana/loki/v3/pkg/logproto"
 	"github.com/grafana/loki/v3/pkg/scratch"
 )
@@ -39,7 +38,7 @@ type builder interface {
 
 // committer allows mocking of certain [kgo.Client] methods in tests.
 type committer interface {
-	CommitRecords(ctx context.Context, records ...*kgo.Record) error
+	Commit(ctx context.Context, offset int64) error
 }
 
 type producer interface {
@@ -52,10 +51,10 @@ type partitionProcessor struct {
 	topic     string
 	partition int32
 	// Processing pipeline
-	records chan *kgo.Record
+	records chan partition.Record
 	// lastRecord contains the last record appended to the builder. It is used
 	// to commit the correct offset after a flush.
-	lastRecord *kgo.Record
+	lastRecord *partition.Record
 	builder    builder
 	decoder    *kafka.Decoder
 	uploader   *uploader.Uploader
@@ -94,14 +93,12 @@ type partitionProcessor struct {
 
 func newPartitionProcessor(
 	ctx context.Context,
-	client *kgo.Client,
+	committer committer,
 	builderCfg logsobj.BuilderConfig,
 	uploaderCfg uploader.Config,
 	metastoreCfg metastore.Config,
 	bucket objstore.Bucket,
 	scratchStore scratch.Store,
-	topic string,
-	partition int32,
 	logger log.Logger,
 	reg prometheus.Registerer,
 	idleFlushTimeout time.Duration,
@@ -112,10 +109,10 @@ func newPartitionProcessor(
 	if err != nil {
 		panic(err)
 	}
-	reg = prometheus.WrapRegistererWith(prometheus.Labels{
-		"topic":     topic,
-		"partition": strconv.Itoa(int(partition)),
-	}, reg)
+	// reg = prometheus.WrapRegistererWith(prometheus.Labels{
+	// 	"topic":     topic,
+	// 	"partition": strconv.Itoa(int(partition)),
+	// }, reg)
 
 	metrics := newPartitionOffsetMetrics()
 	if err := metrics.register(reg); err != nil {
@@ -128,11 +125,9 @@ func newPartitionProcessor(
 	}
 
 	return &partitionProcessor{
-		committer:               client,
-		logger:                  log.With(logger, "partition", partition),
-		topic:                   topic,
-		partition:               partition,
-		records:                 make(chan *kgo.Record, 1000),
+		committer:               committer,
+		logger:                  logger,
+		records:                 make(chan partition.Record, 1000),
 		ctx:                     ctx,
 		cancel:                  cancel,
 		decoder:                 decoder,
@@ -188,7 +183,7 @@ func (p *partitionProcessor) stop() {
 
 // Drops records from the channel if the processor is stopped.
 // Returns false if the processor is stopped, true otherwise.
-func (p *partitionProcessor) Append(records []*kgo.Record) bool {
+func (p *partitionProcessor) Append(records []partition.Record) bool {
 	for _, record := range records {
 		select {
 		// must check per-record in order to not block on a full channel
@@ -241,7 +236,7 @@ func (p *partitionProcessor) emitObjectWrittenEvent(objectPath string) error {
 	return results.FirstErr()
 }
 
-func (p *partitionProcessor) processRecord(record *kgo.Record) {
+func (p *partitionProcessor) processRecord(record partition.Record) {
 	// Update offset metric at the end of processing
 	defer p.metrics.updateOffset(record.Offset)
 
@@ -254,14 +249,8 @@ func (p *partitionProcessor) processRecord(record *kgo.Record) {
 		return
 	}
 
-	tenant := string(record.Key)
-	if !utf8.ValidString(tenant) {
-		// This shouldn't happen, but we catch it here.
-		level.Error(p.logger).Log("msg", "record key is not valid UTF-8")
-		return
-	}
-
-	stream, err := p.decoder.DecodeWithoutLabels(record.Value)
+	tenant := record.TenantID
+	stream, err := p.decoder.DecodeWithoutLabels(record.Content)
 	if err != nil {
 		level.Error(p.logger).Log("msg", "failed to decode record", "err", err)
 		return
@@ -287,7 +276,7 @@ func (p *partitionProcessor) processRecord(record *kgo.Record) {
 		}
 	}
 
-	p.lastRecord = record
+	p.lastRecord = &record
 	p.lastModified = p.clock.Now()
 }
 
@@ -349,7 +338,7 @@ func (p *partitionProcessor) commit() error {
 	backoff.Reset()
 	for backoff.Ongoing() {
 		p.metrics.incCommitsTotal()
-		err := p.committer.CommitRecords(p.ctx, p.lastRecord)
+		err := p.committer.Commit(p.ctx, p.lastRecord.Offset)
 		if err == nil {
 			return nil
 		}

pkg/dataobj/consumer/partition_processor_factory.go

@@ -9,6 +9,7 @@ import (
 	"github.com/twmb/franz-go/pkg/kgo"
 
 	"github.com/grafana/loki/v3/pkg/dataobj/metastore"
+	"github.com/grafana/loki/v3/pkg/kafka/partition"
 	"github.com/grafana/loki/v3/pkg/scratch"
 )
 
@@ -46,31 +47,16 @@ func newPartitionProcessorFactory(
 	}
 }
 
-// New creates a new processor for the per-tenant topic partition.
-//
-// New requires the caller to provide the [kgo.Client] as an argument. This
-// is due to a circular dependency that occurs when creating a [kgo.Client]
-// where the partition event handlers, such as [kgo.OnPartitionsAssigned] and
-// [kgo.OnPartitionsRevoked] must be registered when the client is created.
-// However, the lifecycler cannot be created without the factory, and the
-// factory cannot be created with a [kgo.Client]. This is why New requires a
-// [kgo.Client] as an argument.
-func (f *partitionProcessorFactory) New(
-	ctx context.Context,
-	client *kgo.Client,
-	topic string,
-	partition int32,
-) processor {
+// New returns a new processor for the partition.
+func (f *partitionProcessorFactory) New(committer partition.Committer, logger log.Logger) *partitionProcessor {
 	return newPartitionProcessor(
-		ctx,
-		client,
+		context.TODO(),
+		committer,
 		f.cfg.BuilderConfig,
 		f.cfg.UploaderConfig,
 		f.metastoreCfg,
 		f.bucket,
 		f.scratchStore,
-		topic,
-		partition,
 		f.logger,
 		f.reg,
 		f.cfg.IdleFlushTimeout,