wip

Author: grobinson-grafana

pkg/limits/partition_lifecycler.go

@@ -2,22 +2,33 @@ package limits
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
 	"github.com/twmb/franz-go/pkg/kgo"
+
+	offset "github.com/grafana/loki/v3/pkg/kafka/partition"
 )
 
 // PartitionLifecycler
 type PartitionLifecycler struct {
 	partitionManager *PartitionManager
+	offsetManager    offset.OffsetManager
 	usage            *UsageStore
 	logger           log.Logger
 }
 
 // NewPartitionLifecycler returns a new PartitionLifecycler.
-func NewPartitionLifecycler(partitionManager *PartitionManager, usage *UsageStore, logger log.Logger) *PartitionLifecycler {
+func NewPartitionLifecycler(
+	partitionManager *PartitionManager,
+	offsetManager offset.OffsetManager,
+	usage *UsageStore,
+	logger log.Logger,
+) *PartitionLifecycler {
 	return &PartitionLifecycler{
 		partitionManager: partitionManager,
+		offsetManager: offsetManager,
 		usage:            usage,
 		logger:           log.With(logger, "component", "limits.PartitionLifecycler"),
 	}
@@ -27,6 +38,16 @@ func NewPartitionLifecycler(partitionManager *PartitionManager, usage *UsageStor
 func (l *PartitionLifecycler) Assign(ctx context.Context, _ *kgo.Client, topics map[string][]int32) {
 	for _, partitions := range topics {
 		l.partitionManager.Assign(ctx, partitions)
+		for _, partition := range partitions {
+			if err := l.checkOffsets(ctx, partition); err != nil {
+				level.Error(l.logger).Log(
+					"msg", "failed to check offsets for partition",
+					"partition", partition,
+					"err", err,
+				)
+				l.partitionManager.SetState(partition, PartitionReady)
+			}
+		}
 		return
 	}
 }
@@ -39,3 +60,33 @@ func (l *PartitionLifecycler) Revoke(ctx context.Context, _ *kgo.Client, topics
 		return
 	}
 }
+
+func (l *PartitionLifecycler) checkOffsets(ctx context.Context, partition int32) error {
+	lastProducedOffset, err := l.offsetManager.FetchPartitionOffset(
+		ctx, partition, offset.KafkaEndOffset)
+	if err != nil {
+		return fmt.Errorf("failed to get last produced offset: %w", err)
+	}
+	level.Info(l.logger).Log("lastProducedOffset", lastProducedOffset, "partition", partition)
+	if lastProducedOffset == 0 {
+		// The partition has never produced a record which means there is
+		// nothing to replay and the partition can be marked as ready.
+		l.partitionManager.SetState(partition, PartitionReady)
+		return nil
+	}
+	// Get the last committed offset.
+	lastCommittedOffset, err := l.offsetManager.FetchLastCommittedOffset(ctx, partition)
+	if err != nil {
+		return fmt.Errorf("failed to get last produced offset: %w", err)
+	}
+	level.Info(l.logger).Log("lastProducedOffset", lastProducedOffset, "lastCommittedOffset", lastCommittedOffset, "partition", partition)
+	if lastCommittedOffset == lastProducedOffset {
+		// The last produced record is at least as old as the last committed
+		// offset which means there is nothing to replay either and the
+		// partition can be marked as ready.
+		l.partitionManager.SetState(partition, PartitionReady)
+		return nil
+	}
+	l.partitionManager.SetState(partition, PartitionReplaying)
+	return nil
+}

pkg/limits/partition_manager.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/coder/quartz"
 	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
 )
 
 // PartitionState is the state of a partition in [PartitionManager].
@@ -17,6 +18,20 @@ const (
 	PartitionReady
 )
 
+// String implements the [fmt.Stringer] interface.
+func (s PartitionState) String() string {
+	switch s {
+	case PartitionPending:
+		return "pending"
+	case PartitionReplaying:
+		return "replaying"
+	case PartitionReady:
+		return "ready"
+	default:
+		return "unknown"
+	}
+}
+
 // PartitionManager keeps track of the partitions assigned and for
 // each partition a timestamp of when it was assigned.
 type PartitionManager struct {
@@ -72,6 +87,7 @@ func (m *PartitionManager) SetState(partition int32, state PartitionState) bool
 	defer m.mtx.Unlock()
 	entry, ok := m.partitions[partition]
 	if ok {
+		level.Info(m.logger).Log("msg", "set partition state", "partition", partition, "state", state, "current_state", entry.state)
 		entry.state = state
 		m.partitions[partition] = entry
 	}

pkg/limits/partition_manager_test.go

@@ -173,8 +173,8 @@ func TestPartitionManager_Revoke(t *testing.T) {
 			state:      PartitionPending,
 		},
 	}, m.partitions)
-	// Revoke partitions 2 and 3.
-	m.Revoke(context.Background(), []int32{2, 3})
+	// Remove partitions 2 and 3.
+	m.Remove(context.Background(), []int32{2, 3})
 	require.Equal(t, map[int32]partitionEntry{
 		1: {
 			assignedAt: now,

pkg/limits/playback.go

@@ -46,6 +46,7 @@ func NewPlaybackManager(
 		usage:            usage,
 		readinessCheck:   readinessCheck,
 		logger:           logger,
+		clock: quartz.NewReal(),
 	}
 }
 
@@ -60,55 +61,68 @@ func (m *PlaybackManager) Run(ctx context.Context) error {
 		case <-ctx.Done():
 			return ctx.Err()
 		default:
-			if err := m.processRecords(ctx); err != nil {
+			if err := m.pollFetches(ctx); err != nil {
 				if errors.Is(err, kgo.ErrClientClosed) {
 					return fmt.Errorf("exited: %w", err)
 				}
-				level.Error(m.logger).Log("msg", "failed to process records", "err", err.Error())
+				level.Error(m.logger).Log("msg", "failed to poll fetches", "err", err.Error())
 				b.Wait()
 			}
 		}
 	}
 	return nil
 }
 
-func (m *PlaybackManager) processRecords(ctx context.Context) error {
-	fetches := m.client.PollRecords(ctx, 100)
+func (m *PlaybackManager) pollFetches(ctx context.Context) error {
+	fetches := m.client.PollFetches(ctx)
 	if err := fetches.Err(); err != nil {
 		return err
 	}
-	iter := fetches.RecordIter()
-	for !iter.Done() {
-		if err := m.processRecord(ctx, iter.Next()); err != nil {
-			level.Error(m.logger).Log("msg", "failed to process record", "err", err.Error())
+	fetches.EachPartition(m.processFetchTopicPartition(ctx))
+	return nil
+}
+
+func (m *PlaybackManager) processFetchTopicPartition(ctx context.Context) func(kgo.FetchTopicPartition) {
+	return func(p kgo.FetchTopicPartition) {
+		state, ok := m.partitionManager.GetState(p.Partition)
+		if !ok {
+			level.Warn(m.logger).Log("msg", "discarding records for partition as the partition is not assigned to this client", "partition", p.Partition)
+			return
+		}
+		// Keep a reference to the last processed record. We will use it
+		// for the readiness check if the partition is replaying.
+		var lastProcessed *kgo.Record
+		for _, r := range p.Records {
+			fmt.Println("processing record", r.Timestamp, "partition", p.Partition)
+			if err := m.processRecord(ctx, r); err != nil {
+				level.Error(m.logger).Log("msg", "failed to process record", "err", err.Error())
+			} else {
+				lastProcessed = r
+			}
+		}
+		if lastProcessed != nil {
+			if state == PartitionReplaying {
+				passed, err := m.readinessCheck(lastProcessed)
+				if err != nil {
+					level.Error(m.logger).Log("msg", "failed to run readiness check", "err", err.Error())
+				} else if passed {
+					m.partitionManager.SetState(p.Partition, PartitionReady)
+				}
+			}
+			// If the records were more than one hour ago, commit the offset.
+			// if m.clock.Since(lastProcessed.Timestamp) > time.Second {
+			// 	m.client.CommitRecords(ctx, lastProcessed)
+			// }
 		}
 	}
-	return nil
 }
 
 func (m *PlaybackManager) processRecord(ctx context.Context, r *kgo.Record) error {
-	rec := proto.StreamMetadataRecord{}
-	if err := rec.Unmarshal(r.Value); err != nil {
+	s := proto.StreamMetadataRecord{}
+	if err := s.Unmarshal(r.Value); err != nil {
 		return fmt.Errorf("corrupted record: %w", err)
 	}
-	state, ok := m.partitionManager.GetState(r.Partition)
-	if !ok {
-		return fmt.Errorf("partition is not known to PartitionManager: %d", r.Partition)
-	}
-	// Check if the partition is ready or needs time to replay.
-	if state == PartitionPending || state == PartitionReplaying {
-		ready, err := m.readinessCheck(r)
-		if err != nil {
-			level.Error(m.logger).Log("msg", "failed to check readiness for partition %d: %w", r.Partition, err)
-		} else if ready {
-			m.partitionManager.SetState(r.Partition, PartitionReady)
-		}
-	}
-	m.usage.Update(
-		rec.Tenant,
-		[]*proto.StreamMetadata{rec.Metadata},
-		r.Timestamp.UnixNano(),
-	)
+	m.usage.Update(s.Tenant, []*proto.StreamMetadata{s.Metadata}, r.Timestamp, nil)
 	return nil
 }
 

pkg/limits/service.go

@@ -16,6 +16,7 @@ import (
 	"github.com/twmb/franz-go/pkg/kgo"
 
 	"github.com/grafana/loki/v3/pkg/kafka/client"
+	"github.com/grafana/loki/v3/pkg/kafka/partition"
 	"github.com/grafana/loki/v3/pkg/limits/proto"
 	"github.com/grafana/loki/v3/pkg/util"
 	"github.com/grafana/loki/v3/pkg/util/constants"
@@ -91,10 +92,11 @@ func newMetrics(reg prometheus.Registerer) *metrics {
 type IngestLimits struct {
 	services.Service
 
-	cfg          Config
-	logger       log.Logger
-	clientReader *kgo.Client
-	clientWriter *kgo.Client
+	cfg           Config
+	logger        log.Logger
+	clientReader  *kgo.Client
+	clientWriter  *kgo.Client
+	offsetManager partition.OffsetManager
 
 	lifecycler        *ring.Lifecycler
 	lifecyclerWatcher *services.FailureWatcher
@@ -134,9 +136,9 @@ func NewIngestLimits(cfg Config, lims Limits, logger log.Logger, reg prometheus.
 		cfg:              cfg,
 		logger:           logger,
 		usage:            NewUsageStore(cfg),
+		partitionManager: NewPartitionManager(logger),
 		metrics:          newMetrics(reg),
 		limits:           lims,
-		partitionManager: NewPartitionManager(logger),
 		clock:            quartz.NewReal(),
 	}
 
@@ -161,7 +163,16 @@ func NewIngestLimits(cfg Config, lims Limits, logger log.Logger, reg prometheus.
 	kCfg.AutoCreateTopicEnabled = true
 	kCfg.AutoCreateTopicDefaultPartitions = cfg.NumPartitions
 
-	s.partitionLifecycler = NewPartitionLifecycler(s.partitionManager, s.usage, logger)
+	s.offsetManager, err = partition.NewKafkaOffsetManager(
+		kCfg,
+		"ingest-limits",
+		logger,
+		prometheus.NewRegistry(),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create offset manager: %w", err)
+	}
+	s.partitionLifecycler = NewPartitionLifecycler(s.partitionManager, s.offsetManager, s.usage, logger)
 
 	s.clientReader, err = client.NewReaderClient("ingest-limits-reader", kCfg, logger, reg,
 		kgo.ConsumerGroup(consumerGroup),
@@ -280,6 +291,7 @@ func (s *IngestLimits) starting(ctx context.Context) (err error) {
 func (s *IngestLimits) running(ctx context.Context) error {
 	// Start the eviction goroutine
 	go s.evictOldStreamsPeriodic(ctx)
+	go s.playback.Run(ctx)
 
 	for {
 		select {