feat: Optionally enable kafka histograms for read/write latency (#17089)

Author: aarogoss

docs/sources/shared/configuration.md

@@ -826,6 +826,11 @@ kafka_config:
   # CLI flag: -kafka.max-consumer-lag-at-startup
   [max_consumer_lag_at_startup: <duration> | default = 15s]
 
+  # Enable collection of the following kafka latency histograms: read-wait,
+  # read-timing, write-wait, write-timing
+  # CLI flag: -kafka.enable-kafka-histograms
+  [enable_kafka_histograms: <boolean> | default = false]
+
 dataobj:
   consumer:
     builderconfig:

pkg/blockbuilder/builder/builder.go

@@ -139,7 +139,8 @@ func NewBlockBuilder(
 	logger log.Logger,
 	registerer prometheus.Registerer,
 ) (*BlockBuilder,
-	error) {
+	error,
+) {
 	decoder, err := kafka.NewDecoder()
 	if err != nil {
 		return nil, err
@@ -176,7 +177,7 @@ func (i *BlockBuilder) running(ctx context.Context) error {
 		errgrp.Go(func() error {
 			c, err := client.NewReaderClient(
 				i.kafkaCfg,
-				client.NewReaderClientMetrics(workerID, i.registerer),
+				client.NewReaderClientMetrics(workerID, i.registerer, false),
 				log.With(i.logger, "component", workerID),
 			)
 			if err != nil {
@@ -202,7 +203,6 @@ func (i *BlockBuilder) running(ctx context.Context) error {
 					}
 				}
 			}
-
 		})
 	}
 
@@ -351,7 +351,6 @@ func (i *BlockBuilder) processJob(ctx context.Context, c *kgo.Client, job *types
 		"appender",
 		i.cfg.ConcurrentWriters,
 		func(ctx context.Context) error {
-
 			for {
 				select {
 				case <-ctx.Done():

pkg/dataobj/consumer/service.go

@@ -63,7 +63,7 @@ func New(kafkaCfg kafka.Config, cfg Config, topicPrefix string, bucket objstore.
 		kafkaCfg,
 		partitionRing,
 		groupName,
-		client.NewReaderClientMetrics(groupName, reg),
+		client.NewReaderClientMetrics(groupName, reg, kafkaCfg.EnableKafkaHistograms),
 		logger,
 		kgo.InstanceID(instanceID),
 		kgo.SessionTimeout(3*time.Minute),

pkg/kafka/client/reader_client.go

@@ -43,11 +43,13 @@ func NewReaderClient(kafkaCfg kafka.Config, metrics *kprom.Metrics, logger log.L
 	return client, nil
 }
 
-func NewReaderClientMetrics(component string, reg prometheus.Registerer) *kprom.Metrics {
+func NewReaderClientMetrics(component string, reg prometheus.Registerer, enableKafkaHistograms bool) *kprom.Metrics {
 	return kprom.NewMetrics("loki_ingest_storage_reader",
 		kprom.Registerer(prometheus.WrapRegistererWith(prometheus.Labels{"component": component}, reg)),
 		// Do not export the client ID, because we use it to specify options to the backend.
-		kprom.FetchAndProduceDetail(kprom.Batches, kprom.Records, kprom.CompressedBytes, kprom.UncompressedBytes))
+		kprom.FetchAndProduceDetail(kprom.Batches, kprom.Records, kprom.CompressedBytes, kprom.UncompressedBytes),
+		enableKafkaHistogramMetrics(enableKafkaHistograms),
+	)
 }
 
 // setDefaultNumberOfPartitionsForAutocreatedTopics tries to set num.partitions config option on brokers.
@@ -75,3 +77,24 @@ func setDefaultNumberOfPartitionsForAutocreatedTopics(cfg kafka.Config, cl *kgo.
 
 	level.Info(logger).Log("msg", "configured Kafka-wide default number of partitions for auto-created topics (num.partitions)", "value", cfg.AutoCreateTopicDefaultPartitions)
 }
+
+func enableKafkaHistogramMetrics(enable bool) kprom.Opt {
+	histogramOpts := []kprom.HistogramOpts{}
+	if enable {
+		histogramOpts = append(histogramOpts,
+			kprom.HistogramOpts{
+				Enable:  kprom.ReadTime,
+				Buckets: prometheus.DefBuckets,
+			}, kprom.HistogramOpts{
+				Enable:  kprom.ReadWait,
+				Buckets: prometheus.DefBuckets,
+			}, kprom.HistogramOpts{
+				Enable:  kprom.WriteTime,
+				Buckets: prometheus.DefBuckets,
+			}, kprom.HistogramOpts{
+				Enable:  kprom.WriteWait,
+				Buckets: prometheus.DefBuckets,
+			})
+	}
+	return kprom.HistogramsFromOpts(histogramOpts...)
+}

pkg/kafka/client/writer_client.go

@@ -38,7 +38,9 @@ func NewWriterClient(kafkaCfg kafka.Config, maxInflightProduceRequests int, logg
 	metrics := kprom.NewMetrics(
 		"", // No prefix. We expect the input prometheus.Registered to be wrapped with a prefix.
 		kprom.Registerer(reg),
-		kprom.FetchAndProduceDetail(kprom.Batches, kprom.Records, kprom.CompressedBytes, kprom.UncompressedBytes))
+		kprom.FetchAndProduceDetail(kprom.Batches, kprom.Records, kprom.CompressedBytes, kprom.UncompressedBytes),
+		enableKafkaHistogramMetrics(kafkaCfg.EnableKafkaHistograms),
+	)
 
 	opts := append(
 		commonKafkaClientOptions(kafkaCfg, metrics, logger),

pkg/kafka/config.go

@@ -56,6 +56,8 @@ type Config struct {
 	ProducerMaxBufferedBytes   int64 `yaml:"producer_max_buffered_bytes"`
 
 	MaxConsumerLagAtStartup time.Duration `yaml:"max_consumer_lag_at_startup"`
+
+	EnableKafkaHistograms bool `yaml:"enable_kafka_histograms"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
@@ -85,6 +87,8 @@ func (cfg *Config) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
 
 	consumerLagUsage := fmt.Sprintf("Set -%s to 0 to disable waiting for maximum consumer lag being honored at startup.", prefix+".max-consumer-lag-at-startup")
 	f.DurationVar(&cfg.MaxConsumerLagAtStartup, prefix+".max-consumer-lag-at-startup", 15*time.Second, "The guaranteed maximum lag before a consumer is considered to have caught up reading from a partition at startup, becomes ACTIVE in the hash ring and passes the readiness check. "+consumerLagUsage)
+
+	f.BoolVar(&cfg.EnableKafkaHistograms, prefix+".enable-kafka-histograms", false, "Enable collection of the following kafka latency histograms: read-wait, read-timing, write-wait, write-timing")
 }
 
 func (cfg *Config) Validate() error {

pkg/kafka/partition/offset_manager.go

@@ -92,7 +92,7 @@ func NewKafkaOffsetManager(
 	reg prometheus.Registerer,
 ) (*KafkaOffsetManager, error) {
 	// Create a new Kafka client for the partition manager.
-	clientMetrics := client.NewReaderClientMetrics("partition-manager", reg)
+	clientMetrics := client.NewReaderClientMetrics("partition-manager", reg, cfg.EnableKafkaHistograms)
 	c, err := client.NewReaderClient(
 		cfg,
 		clientMetrics,

pkg/kafka/partition/reader.go

@@ -55,7 +55,7 @@ type ReaderMetrics struct {
 	kprom             *kprom.Metrics
 }
 
-func NewReaderMetrics(r prometheus.Registerer) *ReaderMetrics {
+func NewReaderMetrics(r prometheus.Registerer, enableKafkaHistograms bool) *ReaderMetrics {
 	return &ReaderMetrics{
 		consumptionLag: promauto.With(r).NewHistogramVec(prometheus.HistogramOpts{
 			Name:                            "loki_kafka_reader_consumption_lag_seconds",
@@ -84,7 +84,7 @@ func NewReaderMetrics(r prometheus.Registerer) *ReaderMetrics {
 			Name: "loki_kafka_reader_fetches_total",
 			Help: "Total number of Kafka fetches performed.",
 		}),
-		kprom: client.NewReaderClientMetrics("partition-reader", r),
+		kprom: client.NewReaderClientMetrics("partition-reader", r, enableKafkaHistograms),
 	}
 }
 

pkg/kafka/partition/reader_service.go

@@ -75,7 +75,7 @@ func NewReaderService(
 	logger log.Logger,
 	reg prometheus.Registerer,
 ) (*ReaderService, error) {
-	readerMetrics := NewReaderMetrics(reg)
+	readerMetrics := NewReaderMetrics(reg, kafkaCfg.EnableKafkaHistograms)
 	reader, err := NewKafkaReader(
 		kafkaCfg,
 		partitionID,

pkg/limits/ingest_limits.go

@@ -162,7 +162,7 @@ func NewIngestLimits(cfg Config, logger log.Logger, reg prometheus.Registerer) (
 	s.lifecyclerWatcher = services.NewFailureWatcher()
 	s.lifecyclerWatcher.WatchService(s.lifecycler)
 
-	metrics := client.NewReaderClientMetrics("ingest-limits", reg)
+	metrics := client.NewReaderClientMetrics("ingest-limits", reg, cfg.KafkaConfig.EnableKafkaHistograms)
 
 	// Create a copy of the config to modify the topic
 	kCfg := cfg.KafkaConfig