fix: low watermark with admin api

Signed-off-by: matteo.gazzetta <[email protected]>

Author: amuraru

minion/list_offsets.go

@@ -4,26 +4,49 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"strconv"
 	"time"
 
 	"github.com/twmb/franz-go/pkg/kadm"
 	"go.uber.org/zap"
 )
 
-func (s *Service) ListOffsetsCached(ctx context.Context, timestamp int64) (kadm.ListedOffsets, error) {
+func (s *Service) ListEndOffsetsCached(ctx context.Context) (kadm.ListedOffsets, error) {
+	return s.listOffsetsCached(ctx, "end")
+}
+
+func (s *Service) ListStartOffsetsCached(ctx context.Context) (kadm.ListedOffsets, error) {
+	return s.listOffsetsCached(ctx, "start")
+}
+
+func (s *Service) ListMaxTimestampOffsetsCached(ctx context.Context) (kadm.ListedOffsets, error) {
+	return s.listOffsetsCached(ctx, "maxTimestamp")
+}
+
+func (s *Service) listOffsetsCached(ctx context.Context, offsetType string) (kadm.ListedOffsets, error) {
 	reqId, ok := ctx.Value(RequestIDKey).(string)
 	if !ok || reqId == "" {
 		reqId = "default"
 	}
-	key := "partition-offsets-" + strconv.Itoa(int(timestamp)) + "-" + reqId
+	key := fmt.Sprintf("partition-%s-offsets-%s", offsetType, reqId)
 
 	if cachedRes, exists := s.getCachedItem(key); exists {
 		return cachedRes.(kadm.ListedOffsets), nil
 	}
 
+	var listFunc func(context.Context) (kadm.ListedOffsets, error)
+	switch offsetType {
+	case "end":
+		listFunc = s.ListEndOffsets
+	case "start":
+		listFunc = s.ListStartOffsets
+	case "maxTimestamp":
+		listFunc = s.ListMaxTimestampOffsets
+	default:
+		return nil, fmt.Errorf("invalid offset type: %s", offsetType)
+	}
+
 	res, err, _ := s.requestGroup.Do(key, func() (interface{}, error) {
-		offsets, err := s.ListOffsets(ctx, timestamp)
+		offsets, err := listFunc(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -39,21 +62,41 @@ func (s *Service) ListOffsetsCached(ctx context.Context, timestamp int64) (kadm.
 	return res.(kadm.ListedOffsets), nil
 }
 
-// ListOffsets fetches the low (timestamp: -2) or high water mark (timestamp: -1) for all topic partitions
-func (s *Service) ListOffsets(ctx context.Context, timestamp int64) (kadm.ListedOffsets, error) {
-	listedOffsets, err := s.admClient.ListEndOffsets(ctx)
+// ListEndOffsets fetches the high water mark for all topic partitions.
+func (s *Service) ListEndOffsets(ctx context.Context) (kadm.ListedOffsets, error) {
+	return s.listOffsetsInternal(ctx, s.admClient.ListEndOffsets, "end")
+}
+
+// ListStartOffsets fetches the low water mark for all topic partitions.
+func (s *Service) ListStartOffsets(ctx context.Context) (kadm.ListedOffsets, error) {
+	return s.listOffsetsInternal(ctx, s.admClient.ListStartOffsets, "start")
+}
+
+// ListMaxTimestampOffsets fetches the offsets for the maximum timestamp for all topic partitions.
+// This requires Kafka 3.0+ (see https://issues.apache.org/jira/browse/KAFKA-12541)
+// Uses timestamp -3 which returns the offset for the record with the largest timestamp.
+func (s *Service) ListMaxTimestampOffsets(ctx context.Context) (kadm.ListedOffsets, error) {
+	return s.listOffsetsInternal(ctx, func(ctx context.Context, topics ...string) (kadm.ListedOffsets, error) {
+		return s.admClient.ListOffsetsAfterMilli(ctx, -3, topics...)
+	}, "maxTimestamp")
+}
+
+type listOffsetsFunc func(context.Context, ...string) (kadm.ListedOffsets, error)
+
+func (s *Service) listOffsetsInternal(ctx context.Context, listFunc listOffsetsFunc, offsetType string) (kadm.ListedOffsets, error) {
+	listedOffsets, err := listFunc(ctx)
 	if err != nil {
 		var se *kadm.ShardErrors
 		if !errors.As(err, &se) {
-			return nil, fmt.Errorf("failed to list offsets: %w", err)
+			return nil, fmt.Errorf("failed to list %s offsets: %w", offsetType, err)
 		}
 
 		if se.AllFailed {
-			return nil, fmt.Errorf("failed to list offsets, all shard responses failed: %w", err)
+			return nil, fmt.Errorf("failed to list %s offsets, all shard responses failed: %w", offsetType, err)
 		}
-		s.logger.Info("failed to list offset from some shards", zap.Int("failed_shards", len(se.Errs)))
+		s.logger.Info(fmt.Sprintf("failed to list %s offset from some shards", offsetType), zap.Int("failed_shards", len(se.Errs)))
 		for _, shardErr := range se.Errs {
-			s.logger.Warn("shard error for listing end offsets",
+			s.logger.Warn(fmt.Sprintf("shard error for listing %s offsets", offsetType),
 				zap.Int32("broker_id", shardErr.Broker.NodeID),
 				zap.Error(shardErr.Err))
 		}
@@ -77,12 +120,12 @@ func (s *Service) ListOffsets(ctx context.Context, timestamp int64) (kadm.Listed
 
 	// Print log line for each error type
 	for err, count := range errorCountByErrCode {
-		s.logger.Warn("failed to list some partitions watermarks",
+		s.logger.Warn(fmt.Sprintf("failed to list some partitions %s watermarks", offsetType),
 			zap.Error(err),
 			zap.Int("error_count", count))
 	}
 	if len(errorCountByTopic) > 0 {
-		s.logger.Warn("some topics had one or more partitions whose watermarks could not be fetched from Kafka",
+		s.logger.Warn(fmt.Sprintf("some topics had one or more partitions whose %s watermarks could not be fetched from Kafka", offsetType),
 			zap.Int("topics_with_errors", len(errorCountByTopic)))
 	}
 

prometheus/collect_consumer_group_lags.go

@@ -26,13 +26,13 @@ func (e *Exporter) collectConsumerGroupLags(ctx context.Context, ch chan<- prome
 
 	// Low Watermarks (at the moment they are not needed at all, they could be used to calculate the lag on partitions
 	// that don't have any active offsets)
-	lowWaterMarks, err := e.minionSvc.ListOffsetsCached(ctx, -2)
+	lowWaterMarks, err := e.minionSvc.ListStartOffsetsCached(ctx)
 	if err != nil {
 		e.logger.Error("failed to fetch low water marks", zap.Error(err))
 		return false
 	}
 	// High Watermarks
-	highWaterMarks, err := e.minionSvc.ListOffsetsCached(ctx, -1)
+	highWaterMarks, err := e.minionSvc.ListEndOffsetsCached(ctx)
 	if err != nil {
 		e.logger.Error("failed to fetch low water marks", zap.Error(err))
 		return false

prometheus/collect_topic_partition_offsets.go

@@ -17,13 +17,13 @@ func (e *Exporter) collectTopicPartitionOffsets(ctx context.Context, ch chan<- p
 	isOk := true
 
 	// Low Watermarks
-	lowWaterMarks, err := e.minionSvc.ListOffsetsCached(ctx, -2)
+	lowWaterMarks, err := e.minionSvc.ListStartOffsetsCached(ctx)
 	if err != nil {
 		e.logger.Error("failed to fetch low water marks", zap.Error(err))
 		return false
 	}
 	// High Watermarks
-	highWaterMarks, err := e.minionSvc.ListOffsetsCached(ctx, -1)
+	highWaterMarks, err := e.minionSvc.ListEndOffsetsCached(ctx)
 	if err != nil {
 		e.logger.Error("failed to fetch low water marks", zap.Error(err))
 		return false
@@ -32,7 +32,7 @@ func (e *Exporter) collectTopicPartitionOffsets(ctx context.Context, ch chan<- p
 	// Highest Timestamp Offsets
 	// NB: this requires Kafka Brokers 3.0+ (see https://issues.apache.org/jira/browse/KAFKA-12541)
 	// In older versions this is returning the timestamp of the low watermarks (earliest offset)
-	maxTimestampOffsets, err := e.minionSvc.ListOffsetsCached(ctx, -3)
+	maxTimestampOffsets, err := e.minionSvc.ListMaxTimestampOffsetsCached(ctx)
 	if err != nil {
 		e.logger.Error("failed to fetch offsets for max timestamp", zap.Error(err))
 		return false