feat: support replaying records for newly assigned partitions

This commit supports replaying records for newly assigned partitions
so an instance can rebuild it's in-memory usage data before serving
queries.

Author: grobinson-grafana

pkg/blockbuilder/scheduler/scheduler_test.go

@@ -29,6 +29,9 @@ func (m *mockOffsetManager) ConsumerGroup() string { return m.consumerGroup }
 func (m *mockOffsetManager) GroupLag(_ context.Context, _ int64) (map[int32]partition.Lag, error) {
 	return nil, nil
 }
+func (m *mockOffsetManager) NextOffset(_ context.Context, _ int32, _ time.Time) (int64, error) {
+	return 0, nil
+}
 func (m *mockOffsetManager) FetchLastCommittedOffset(_ context.Context, _ int32) (int64, error) {
 	return 0, nil
 }

pkg/kafka/partition/offset_manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -72,6 +73,7 @@ type OffsetManager interface {
 	GroupLag(ctx context.Context, fallbackOffsetMillis int64) (map[int32]Lag, error)
 	FetchLastCommittedOffset(ctx context.Context, partition int32) (int64, error)
 	FetchPartitionOffset(ctx context.Context, partition int32, position SpecialOffset) (int64, error)
+	NextOffset(ctx context.Context, partition int32, t time.Time) (int64, error)
 	Commit(ctx context.Context, partition int32, offset int64) error
 }
 
@@ -130,6 +132,31 @@ func (r *KafkaOffsetManager) ConsumerGroup() string {
 	return r.cfg.GetConsumerGroup(r.instanceID)
 }
 
+// NextOffset returns the first offset after the timestamp t. If the partition
+// does not have an offset after t, it returns the current end offset.
+func (r *KafkaOffsetManager) NextOffset(ctx context.Context, partition int32, t time.Time) (int64, error) {
+	resp, err := r.adminClient.ListOffsetsAfterMilli(ctx, t.UnixMilli(), r.cfg.Topic)
+	if err != nil {
+		return 0, err
+	}
+	// If a topic does not exist, a special -1 partition for each non-existing
+	// topic is added to the response.
+	partitions := resp[r.cfg.Topic]
+	if special, ok := partitions[-1]; ok {
+		return 0, special.Err
+	}
+	// If a partition does not exist, it will be missing.
+	listed, ok := partitions[partition]
+	if !ok {
+		return 0, fmt.Errorf("unknown partition %d", partition)
+	}
+	// Err is non-nil if the partition has a load error.
+	if listed.Err != nil {
+		return 0, listed.Err
+	}
+	return listed.Offset, nil
+}
+
 // FetchLastCommittedOffset retrieves the last committed offset for this partition
 func (r *KafkaOffsetManager) FetchLastCommittedOffset(ctx context.Context, partitionID int32) (int64, error) {
 	req := kmsg.NewPtrOffsetFetchRequest()
@@ -218,7 +245,6 @@ func (r *KafkaOffsetManager) FetchPartitionOffset(ctx context.Context, partition
 	if err := kerr.ErrorForCode(partition.ErrorCode); err != nil {
 		return 0, err
 	}
-
 	return partition.Offset, nil
 }
 

pkg/limits/http_test.go

@@ -42,8 +42,10 @@ func TestIngestLimits_ServeHTTP(t *testing.T) {
 		},
 		logger: log.NewNopLogger(),
 		partitionManager: &PartitionManager{
-			partitions: map[int32]int64{
-				0: time.Now().UnixNano(),
+			partitions: map[int32]partitionEntry{
+				0: {
+					assignedAt: time.Now().UnixNano(),
+				},
 			},
 		},
 	}

pkg/limits/mock_test.go

@@ -6,8 +6,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
-
-	"github.com/grafana/loki/v3/pkg/limits/proto"
+	"github.com/twmb/franz-go/pkg/kgo"
 )
 
 type MockLimits struct {
@@ -27,24 +26,26 @@ func (m *MockLimits) IngestionBurstSizeBytes(_ string) int {
 	return 1000
 }
 
-type mockWAL struct {
-	t                    *testing.T
-	NumAppendsTotal      int
-	ExpectedAppendsTotal int
-	mtx                  sync.Mutex
-}
+// mockKafka mocks a [kgo.Client].
+type mockKafka struct {
+	t   *testing.T
+	mtx sync.Mutex
 
-func (m *mockWAL) Append(_ context.Context, _ string, _ *proto.StreamMetadata) error {
-	m.mtx.Lock()
-	defer m.mtx.Unlock()
-	m.NumAppendsTotal++
-	return nil
+	expectedNumRecords int
+	numRecords         int
 }
 
-func (m *mockWAL) Close() error {
-	return nil
+func (k *mockKafka) Produce(
+	_ context.Context,
+	r *kgo.Record,
+	promise func(*kgo.Record, error),
+) {
+	k.mtx.Lock()
+	defer k.mtx.Unlock()
+	k.numRecords++
+	promise(r, nil)
 }
 
-func (m *mockWAL) AssertAppendsTotal() {
-	require.Equal(m.t, m.ExpectedAppendsTotal, m.NumAppendsTotal)
+func (k *mockKafka) AssertExpectedNumRecords() {
+	require.Equal(k.t, k.expectedNumRecords, k.numRecords)
 }

pkg/limits/partition_lifecycler.go

@@ -2,24 +2,39 @@ package limits
 
 import (
 	"context"
+	"fmt"
+	"time"
 
 	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
 	"github.com/twmb/franz-go/pkg/kgo"
+
+	kafka_partition "github.com/grafana/loki/v3/pkg/kafka/partition"
 )
 
 // PartitionLifecycler manages assignment and revocation of partitions.
 type PartitionLifecycler struct {
+	cfg              Config
 	partitionManager *PartitionManager
+	offsetManager    kafka_partition.OffsetManager
 	usage            *UsageStore
 	logger           log.Logger
 }
 
 // NewPartitionLifecycler returns a new PartitionLifecycler.
-func NewPartitionLifecycler(partitionManager *PartitionManager, usage *UsageStore, logger log.Logger) *PartitionLifecycler {
+func NewPartitionLifecycler(
+	cfg Config,
+	partitionManager *PartitionManager,
+	offsetManager kafka_partition.OffsetManager,
+	usage *UsageStore,
+	logger log.Logger,
+) *PartitionLifecycler {
 	return &PartitionLifecycler{
+		cfg:              cfg,
 		partitionManager: partitionManager,
+		offsetManager:    offsetManager,
 		usage:            usage,
-		logger:           log.With(logger, "component", "limits.PartitionLifecycler"),
+		logger:           logger,
 	}
 }
 
@@ -29,6 +44,16 @@ func (l *PartitionLifecycler) Assign(ctx context.Context, _ *kgo.Client, topics
 	// TODO(grobinson): Figure out what to do if this is not the case.
 	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, partition is ready",
+					"partition", partition,
+					"err", err,
+				)
+				l.partitionManager.SetReady(partition)
+			}
+		}
 		return
 	}
 }
@@ -43,3 +68,55 @@ func (l *PartitionLifecycler) Revoke(ctx context.Context, _ *kgo.Client, topics
 		return
 	}
 }
+
+func (l *PartitionLifecycler) checkOffsets(ctx context.Context, partition int32) error {
+	logger := log.With(l.logger, "partition", partition)
+	// Get the start offset for the partition. This can be greater than zero
+	// if a retention period has deleted old records.
+	startOffset, err := l.offsetManager.FetchPartitionOffset(
+		ctx, partition, kafka_partition.KafkaStartOffset)
+	if err != nil {
+		return fmt.Errorf("failed to get last produced offset: %w", err)
+	}
+	// The last produced offset is the next offset after the last produced
+	// record. For example, if a partition contains 1 record, then the last
+	// produced offset is 1. However, the offset of the last produced record
+	// is 0, as offsets start from 0.
+	lastProducedOffset, err := l.offsetManager.FetchPartitionOffset(
+		ctx, partition, kafka_partition.KafkaEndOffset)
+	if err != nil {
+		return fmt.Errorf("failed to get last produced offset: %w", err)
+	}
+	// Get the first offset produced within the window. This can be the same
+	// offset as the last produced offset if no records have been produced
+	// within that time.
+	nextOffset, err := l.offsetManager.NextOffset(ctx, partition, time.Now().Add(-l.cfg.WindowSize))
+	if err != nil {
+		return fmt.Errorf("failed to get next offset: %w", err)
+	}
+	level.Debug(logger).Log(
+		"msg", "fetched offsets",
+		"start_offset", startOffset,
+		"last_produced_offset", lastProducedOffset,
+		"next_offset", nextOffset,
+	)
+	if startOffset >= lastProducedOffset {
+		// The partition has no records. This happens when either the
+		// partition has never produced a record, or all records that have
+		// been produced have been deleted due to the retention period.
+		level.Debug(logger).Log("msg", "no records in partition, partition is ready")
+		l.partitionManager.SetReady(partition)
+		return nil
+	}
+	if nextOffset == lastProducedOffset {
+		level.Debug(logger).Log("msg", "no records within window size, partition is ready")
+		l.partitionManager.SetReady(partition)
+		return nil
+	}
+	// Since we want to fetch all records up to and including the last
+	// produced record, we must fetch all records up to and including the
+	// last produced offset - 1.
+	level.Debug(logger).Log("msg", "partition is replaying")
+	l.partitionManager.SetReplaying(partition, lastProducedOffset-1)
+	return nil
+}