feat: discard records for our own zone

Author: grobinson-grafana

pkg/limits/playback_manager.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"strconv"
 	"time"
 
 	"github.com/coder/quartz"
@@ -33,11 +32,15 @@ type PlaybackManager struct {
 	// readinessCheck checks if a waiting or replaying partition can be
 	// switched to ready.
 	readinessCheck PartitionReadinessCheck
-	logger         log.Logger
+	// zone is used to discard our own records.
+	zone   string
+	logger log.Logger
 
 	// Metrics.
-	recordsProcessed *prometheus.CounterVec
 	lag              prometheus.Histogram
+	recordsFetched   prometheus.Counter
+	recordsDiscarded prometheus.Counter
+	recordsInvalid   prometheus.Counter
 
 	// Used for tests.
 	clock quartz.Clock
@@ -49,6 +52,7 @@ func NewPlaybackManager(
 	partitionManager *PartitionManager,
 	usage *UsageStore,
 	readinessCheck PartitionReadinessCheck,
+	zone string,
 	logger log.Logger,
 	reg prometheus.Registerer,
 ) *PlaybackManager {
@@ -57,15 +61,9 @@ func NewPlaybackManager(
 		partitionManager: partitionManager,
 		usage:            usage,
 		readinessCheck:   readinessCheck,
+		zone:             zone,
 		logger:           logger,
 		clock:            quartz.NewReal(),
-		recordsProcessed: promauto.With(reg).NewCounterVec(
-			prometheus.CounterOpts{
-				Name: "loki_ingest_limits_records_processed_total",
-				Help: "The total number of records processed.",
-			},
-			[]string{"partition"},
-		),
 		lag: promauto.With(reg).NewHistogram(
 			prometheus.HistogramOpts{
 				Name:                            "loki_ingest_limits_lag_seconds",
@@ -76,6 +74,24 @@ func NewPlaybackManager(
 				Buckets:                         prometheus.ExponentialBuckets(0.125, 2, 18),
 			},
 		),
+		recordsFetched: promauto.With(reg).NewCounter(
+			prometheus.CounterOpts{
+				Name: "loki_ingest_limits_records_fetched_total",
+				Help: "The total number of records fetched.",
+			},
+		),
+		recordsDiscarded: promauto.With(reg).NewCounter(
+			prometheus.CounterOpts{
+				Name: "loki_ingest_limits_records_discarded_total",
+				Help: "The total number of records discarded.",
+			},
+		),
+		recordsInvalid: promauto.With(reg).NewCounter(
+			prometheus.CounterOpts{
+				Name: "loki_ingest_limits_records_invalid_total",
+				Help: "The total number of invalid records.",
+			},
+		),
 	}
 }
 
@@ -121,26 +137,26 @@ func (m *PlaybackManager) processFetchTopicPartition(ctx context.Context) func(k
 			return
 		}
 		logger := log.With(m.logger, "partition", p.Partition)
+		m.recordsFetched.Add(float64(len(p.Records)))
+		// We need the state of the partition so we can discard any records
+		// that we produced (unless replaying) and mark a replaying partition
+		// as ready once it has finished replaying.
 		state, ok := m.partitionManager.GetState(p.Partition)
 		if !ok {
+			m.recordsDiscarded.Add(float64(len(p.Records)))
 			level.Warn(logger).Log("msg", "discarding records for partition as the partition is not assigned to this client")
 			return
 		}
-		if state == PartitionReplaying {
-			// TODO(grobinson): For now we just consume records when replaying
-			// a partition. In a future commit this will be moved outside of
-			// this check, and records will be consumed both when replaying
-			// newly assigned partitions and merging records from other zones.
-			for _, r := range p.Records {
-				if err := m.processRecord(ctx, r); err != nil {
-					level.Error(logger).Log("msg", "failed to process record", "err", err.Error())
-				}
+		for _, r := range p.Records {
+			if err := m.processRecord(ctx, state, r); err != nil {
+				level.Error(logger).Log("msg", "failed to process record", "err", err.Error())
 			}
-			m.recordsProcessed.
-				WithLabelValues(strconv.FormatInt(int64(p.Partition), 10)).
-				Add(float64(len(p.Records)))
-			m.lag.Observe(m.clock.Since(p.Records[len(p.Records)-1].Timestamp).Seconds())
-			passed, err := m.readinessCheck(p.Partition, p.Records[len(p.Records)-1])
+		}
+		// Get the last record (has the latest offset and timestamp).
+		lastRecord := p.Records[len(p.Records)-1]
+		m.lag.Observe(m.clock.Since(lastRecord.Timestamp).Seconds())
+		if state == PartitionReplaying {
+			passed, err := m.readinessCheck(p.Partition, lastRecord)
 			if err != nil {
 				level.Error(logger).Log("msg", "failed to run readiness check", "err", err.Error())
 			} else if passed {
@@ -151,11 +167,17 @@ func (m *PlaybackManager) processFetchTopicPartition(ctx context.Context) func(k
 	}
 }
 
-func (m *PlaybackManager) processRecord(_ context.Context, r *kgo.Record) error {
+func (m *PlaybackManager) processRecord(_ context.Context, state PartitionState, r *kgo.Record) error {
 	s := proto.StreamMetadataRecord{}
 	if err := s.Unmarshal(r.Value); err != nil {
+		m.recordsInvalid.Inc()
 		return fmt.Errorf("corrupted record: %w", err)
 	}
+	if state == PartitionReady && m.zone == s.Zone {
+		// Discard our own records so we don't count the same streams twice.
+		m.recordsDiscarded.Inc()
+		return nil
+	}
 	m.usage.Update(s.Tenant, []*proto.StreamMetadata{s.Metadata}, r.Timestamp, nil)
 	return nil
 }

pkg/limits/sender.go

@@ -26,18 +26,20 @@ type Sender struct {
 	// set in the client.
 	topic      string
 	partitions int
+	zone       string
 	logger     log.Logger
 
 	produced       prometheus.Counter
 	producedFailed prometheus.Counter
 }
 
 // NewSender returns a new Sender.
-func NewSender(producer Producer, topic string, partitions int, logger log.Logger, reg prometheus.Registerer) *Sender {
+func NewSender(producer Producer, topic string, partitions int, zone string, logger log.Logger, reg prometheus.Registerer) *Sender {
 	return &Sender{
 		producer:   producer,
 		topic:      topic,
 		partitions: partitions,
+		zone:       zone,
 		logger:     logger,
 		produced: promauto.With(reg).NewCounter(
 			prometheus.CounterOpts{
@@ -59,6 +61,7 @@ func NewSender(producer Producer, topic string, partitions int, logger log.Logge
 // complete.
 func (s *Sender) Produce(ctx context.Context, tenant string, metadata *proto.StreamMetadata) error {
 	v := proto.StreamMetadataRecord{
+		Zone:     s.zone,
 		Tenant:   tenant,
 		Metadata: metadata,
 	}

pkg/limits/service.go

@@ -186,10 +186,18 @@ func NewIngestLimits(cfg Config, lims Limits, logger log.Logger, reg prometheus.
 		s.partitionManager,
 		s.usage,
 		NewOffsetReadinessCheck(s.partitionManager),
+		cfg.LifecyclerConfig.Zone,
+		logger,
+		reg,
+	)
+	s.sender = NewSender(
+		s.clientWriter,
+		kCfg.Topic,
+		s.cfg.NumPartitions,
+		cfg.LifecyclerConfig.Zone,
 		logger,
 		reg,
 	)
-	s.sender = NewSender(s.clientWriter, kCfg.Topic, s.cfg.NumPartitions, logger, reg)
 
 	s.Service = services.NewBasicService(s.starting, s.running, s.stopping)
 	return s, nil

pkg/limits/service_test.go

@@ -334,7 +334,7 @@ func TestIngestLimits_ExceedsLimits(t *testing.T) {
 				usage:            tt.usage,
 				partitionManager: NewPartitionManager(),
 				clock:            clock,
-				sender:           NewSender(&kafkaClient, "test", tt.numPartitions, log.NewNopLogger(), reg),
+				sender:           NewSender(&kafkaClient, "test", tt.numPartitions, "", log.NewNopLogger(), reg),
 			}
 
 			// Assign the Partition IDs.
@@ -423,7 +423,7 @@ func TestIngestLimits_ExceedsLimits_Concurrent(t *testing.T) {
 		metrics:          newMetrics(reg),
 		limits:           limits,
 		clock:            clock,
-		sender:           NewSender(&kafkaClient, "test", 1, log.NewNopLogger(), reg),
+		sender:           NewSender(&kafkaClient, "test", 1, "", log.NewNopLogger(), reg),
 	}
 
 	// Assign the Partition IDs.