Collect consumer metrics

zio-kafka/src/main/scala/zio/kafka/consumer/ConsumerSettings.scala

@@ -30,7 +30,8 @@ final case class ConsumerSettings(
   restartStreamOnRebalancing: Boolean = false,
   rebalanceSafeCommits: Boolean = false,
   maxRebalanceDuration: Option[Duration] = None,
-  fetchStrategy: FetchStrategy = QueueSizeBasedFetchStrategy()
+  fetchStrategy: FetchStrategy = QueueSizeBasedFetchStrategy(),
+  metricsConsumerId: Option[String] = None
 ) {
 
   /**
@@ -278,6 +279,21 @@ final case class ConsumerSettings(
    */
   def withFetchStrategy(fetchStrategy: FetchStrategy): ConsumerSettings =
     copy(fetchStrategy = fetchStrategy)
+
+  /**
+   * @param consumerId
+   *   The value given to the metrics label `consumer_id` for all metrics. When this value is not set, the consumer
+   *   group id is used, when no group id is set, it defaults to a random value that is generated when the consumer is
+   *   created.
+   */
+  def withMetricsConsumerId(consumerId: String): ConsumerSettings =
+    copy(metricsConsumerId = Some(consumerId))
+
+  def derivedMetricsConsumerId: String =
+    metricsConsumerId
+      .orElse(properties.get(ConsumerConfig.GROUP_ID_CONFIG).map(_.toString))
+      .getOrElse(Seq.fill(6)(scala.util.Random.nextInt(16).toHexString).mkString)
+
 }
 
 object ConsumerSettings {

zio-kafka/src/main/scala/zio/kafka/consumer/internal/ConsumerMetrics.scala

@@ -0,0 +1,66 @@
+package zio.kafka.consumer.internal
+
+import zio.metrics.MetricKeyType.Histogram
+import zio.metrics._
+import zio.{ Chunk, UIO, ZIO }
+
+final case class ConsumerMetrics(metricsConsumerId: String) {
+
+  // Chunk(0,1,3,8,21,55,149,404,1097,2981)
+  private val streamCountBoundaries: Histogram.Boundaries =
+    MetricKeyType.Histogram.Boundaries.fromChunk(Chunk(0.0) ++ Chunk.iterate(1.0, 9)(_ * Math.E).map(Math.ceil))
+
+  // Chunk(0,100,272,739,2009,5460,14842,40343,109664,298096)
+  private val streamSizeBoundaries: Histogram.Boundaries =
+    MetricKeyType.Histogram.Boundaries.fromChunk(Chunk(0.0) ++ Chunk.iterate(100.0, 9)(_ * Math.E).map(Math.ceil))
+
+  private val pendingRequestsHistogram =
+    Metric
+      .histogram(
+        "ziokafka_consumer_pending_requests",
+        "The number of partition streams that are awaiting new records.",
+        streamCountBoundaries
+      )
+      .tagged("consumer_id", metricsConsumerId)
+
+  private val pendingCommitsHistogram =
+    Metric
+      .histogram(
+        "ziokafka_consumer_pending_commits",
+        "The number of commits that are awaiting completion.",
+        streamCountBoundaries
+      )
+      .tagged("consumer_id", metricsConsumerId)
+
+  private val queueSizeHistogram =
+    Metric
+      .histogram(
+        "ziokafka_consumer_queue_size",
+        "The number of records queued per partition.",
+        streamSizeBoundaries
+      )
+      .tagged("consumer_id", metricsConsumerId)
+
+  private val allQueueSizeHistogram =
+    Metric
+      .histogram(
+        "ziokafka_consumer_all_queue_size",
+        "The number of records queued in the consumer (all partitions).",
+        streamSizeBoundaries
+      )
+      .tagged("consumer_id", metricsConsumerId)
+
+  def observeMetrics(state: Runloop.State): UIO[Unit] =
+    ZIO
+      .when(state.subscriptionState.isSubscribed) {
+        for {
+          queueSizes <- ZIO.foreach(state.assignedStreams)(_.queueSize.map(_.toDouble))
+          _          <- ZIO.foreachDiscard(queueSizes)(qs => queueSizeHistogram.update(qs))
+          _          <- allQueueSizeHistogram.update(queueSizes.sum)
+          _          <- ZIO.succeed(state.pendingRequests.size.toDouble) @@ pendingRequestsHistogram
+          _          <- ZIO.succeed(state.pendingCommits.size.toDouble) @@ pendingCommitsHistogram
+        } yield ()
+      }
+      .unit
+
+}

zio-kafka/src/main/scala/zio/kafka/consumer/internal/Runloop.scala

@@ -22,6 +22,7 @@ import scala.jdk.CollectionConverters._
 
 //noinspection SimplifyWhenInspection,SimplifyUnlessInspection
 private[consumer] final class Runloop private (
+  metricsConsumerId: String,
   sameThreadRuntime: Runtime[Any],
   hasGroupId: Boolean,
   consumer: ConsumerAccess,
@@ -43,6 +44,8 @@ private[consumer] final class Runloop private (
   fetchStrategy: FetchStrategy
 ) {
 
+  private val consumerMetrics = ConsumerMetrics(metricsConsumerId)
+
   private def newPartitionStream(tp: TopicPartition): UIO[PartitionStreamControl] =
     PartitionStreamControl.newPartitionStream(tp, commandQueue, diagnostics, maxPollInterval)
 
@@ -683,6 +686,7 @@ private[consumer] final class Runloop private (
       .takeWhile(_ != RunloopCommand.StopRunloop)
       .runFoldChunksDiscardZIO(initialState) { (state, commands) =>
         for {
+          _              <- consumerMetrics.observeMetrics(state).fork
           commitCommands <- commitQueue.takeAll
           _ <- ZIO.logDebug(
                  s"Processing ${commitCommands.size} commits," +
@@ -785,6 +789,7 @@ object Runloop {
   }
 
   private[consumer] def make(
+    metricsConsumerId: String,
     hasGroupId: Boolean,
     consumer: ConsumerAccess,
     pollTimeout: Duration,
@@ -809,6 +814,7 @@ object Runloop {
       committedOffsetsRef <- Ref.make(CommitOffsets.empty)
       sameThreadRuntime   <- ZIO.runtime[Any].provideLayer(SameThreadRuntimeLayer)
       runloop = new Runloop(
+                  metricsConsumerId = metricsConsumerId,
                   sameThreadRuntime = sameThreadRuntime,
                   hasGroupId = hasGroupId,
                   consumer = consumer,
@@ -844,7 +850,7 @@ object Runloop {
            )
     } yield runloop
 
-  private final case class State(
+  private[internal] final case class State(
     pendingRequests: Chunk[RunloopCommand.Request],
     pendingCommits: Chunk[Runloop.Commit],
     assignedStreams: Chunk[PartitionStreamControl],

zio-kafka/src/main/scala/zio/kafka/consumer/internal/RunloopAccess.scala

@@ -89,6 +89,7 @@ private[consumer] object RunloopAccess {
       runloopStateRef <- Ref.Synchronized.make[RunloopState](RunloopState.NotStarted)
       makeRunloop = Runloop
                       .make(
+                        metricsConsumerId = settings.derivedMetricsConsumerId,
                         hasGroupId = settings.hasGroupId,
                         consumer = consumerAccess,
                         pollTimeout = settings.pollTimeout,