Collect consumer metrics

zio-kafka-test/src/test/scala/zio/kafka/consumer/internal/RunloopCommitOffsetsSpec.scala

@@ -16,39 +16,57 @@ object RunloopCommitOffsetsSpec extends ZIOSpecDefault {
   override def spec: Spec[TestEnvironment with Scope, Any] =
     suite("Runloop.CommitOffsets spec")(
       test("addCommits adds to empty CommitOffsets") {
-        val s1 = Runloop.CommitOffsets(Map.empty)
-        val s2 = s1.addCommits(Chunk(makeCommit(Map(tp10 -> 10))))
-        assertTrue(s2.offsets == Map(tp10 -> 10L))
+        val s1        = Runloop.CommitOffsets(Map.empty)
+        val (inc, s2) = s1.addCommits(Chunk(makeCommit(Map(tp10 -> 10))))
+        assertTrue(
+          inc == 0,
+          s2.offsets == Map(tp10 -> 10L)
+        )
       },
       test("addCommits updates offset when it is higher") {
-        val s1 = Runloop.CommitOffsets(Map(tp10 -> 5L))
-        val s2 = s1.addCommits(Chunk(makeCommit(Map(tp10 -> 10))))
-        assertTrue(s2.offsets == Map(tp10 -> 10L))
+        val s1        = Runloop.CommitOffsets(Map(tp10 -> 4L))
+        val (inc, s2) = s1.addCommits(Chunk(makeCommit(Map(tp10 -> 10))))
+        assertTrue(
+          inc == 10 - 4,
+          s2.offsets == Map(tp10 -> 10L)
+        )
       },
       test("addCommits ignores an offset when it is lower") {
-        val s1 = Runloop.CommitOffsets(Map(tp10 -> 10L))
-        val s2 = s1.addCommits(Chunk(makeCommit(Map(tp10 -> 5))))
-        assertTrue(s2.offsets == Map(tp10 -> 10L))
+        val s1        = Runloop.CommitOffsets(Map(tp10 -> 10L))
+        val (inc, s2) = s1.addCommits(Chunk(makeCommit(Map(tp10 -> 5))))
+        assertTrue(
+          inc == 0,
+          s2.offsets == Map(tp10 -> 10L)
+        )
       },
       test("addCommits keeps unrelated partitions") {
-        val s1 = Runloop.CommitOffsets(Map(tp10 -> 10L))
-        val s2 = s1.addCommits(Chunk(makeCommit(Map(tp11 -> 11))))
-        assertTrue(s2.offsets == Map(tp10 -> 10L, tp11 -> 11L))
+        val s1        = Runloop.CommitOffsets(Map(tp10 -> 10L))
+        val (inc, s2) = s1.addCommits(Chunk(makeCommit(Map(tp11 -> 11))))
+        assertTrue(
+          inc == 0,
+          s2.offsets == Map(tp10 -> 10L, tp11 -> 11L)
+        )
       },
       test("addCommits does it all at once") {
-        val s1 = Runloop.CommitOffsets(Map(tp10 -> 10L, tp20 -> 205L, tp21 -> 210L, tp22 -> 220L))
-        val s2 = s1.addCommits(Chunk(makeCommit(Map(tp11 -> 11, tp20 -> 206L, tp21 -> 209L, tp22 -> 220L))))
-        assertTrue(s2.offsets == Map(tp10 -> 10L, tp11 -> 11L, tp20 -> 206L, tp21 -> 210L, tp22 -> 220L))
+        val s1        = Runloop.CommitOffsets(Map(tp10 -> 10L, tp20 -> 205L, tp21 -> 210L, tp22 -> 220L))
+        val (inc, s2) = s1.addCommits(Chunk(makeCommit(Map(tp11 -> 11, tp20 -> 206L, tp21 -> 209L, tp22 -> 220L))))
+        assertTrue(
+          inc == /* tp10 */ 0 + /* tp11 */ 0 + /* tp20 */ 1 + /* tp21 */ 0 + /* tp22 */ 0,
+          s2.offsets == Map(tp10 -> 10L, tp11 -> 11L, tp20 -> 206L, tp21 -> 210L, tp22 -> 220L)
+        )
       },
       test("addCommits adds multiple commits") {
         val s1 = Runloop.CommitOffsets(Map(tp10 -> 10L, tp20 -> 200L, tp21 -> 210L, tp22 -> 220L))
-        val s2 = s1.addCommits(
+        val (inc, s2) = s1.addCommits(
           Chunk(
             makeCommit(Map(tp11 -> 11, tp20 -> 199L, tp21 -> 211L, tp22 -> 219L)),
             makeCommit(Map(tp20 -> 198L, tp21 -> 209L, tp22 -> 221L))
           )
         )
-        assertTrue(s2.offsets == Map(tp10 -> 10L, tp11 -> 11L, tp20 -> 200L, tp21 -> 211L, tp22 -> 221L))
+        assertTrue(
+          inc == /* tp10 */ 0 + /* tp11 */ 0 + /* tp20 */ 0 + /* tp21 */ 1 + /* tp22 */ 1,
+          s2.offsets == Map(tp10 -> 10L, tp11 -> 11L, tp20 -> 200L, tp21 -> 211L, tp22 -> 221L)
+        )
       },
       test("keepPartitions removes some partitions") {
         val s1 = Runloop.CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
@@ -80,6 +98,6 @@ object RunloopCommitOffsetsSpec extends ZIOSpecDefault {
   private def makeCommit(offsets: Map[TopicPartition, Long]): Runloop.Commit = {
     val o = offsets.map { case (tp, offset) => tp -> new OffsetAndMetadata(offset) }
     val p = Unsafe.unsafe(implicit unsafe => Promise.unsafe.make[Throwable, Unit](FiberId.None))
-    Runloop.Commit(o, p)
+    Runloop.Commit(0L, o, p)
   }
 }

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

@@ -5,6 +5,7 @@ import zio._
 import zio.kafka.consumer.Consumer.OffsetRetrieval
 import zio.kafka.consumer.fetch.{ FetchStrategy, QueueSizeBasedFetchStrategy }
 import zio.kafka.security.KafkaCredentialStore
+import zio.metrics.MetricLabel
 
 /**
  * Settings for the consumer.
@@ -30,7 +31,9 @@ final case class ConsumerSettings(
   restartStreamOnRebalancing: Boolean = false,
   rebalanceSafeCommits: Boolean = false,
   maxRebalanceDuration: Option[Duration] = None,
-  fetchStrategy: FetchStrategy = QueueSizeBasedFetchStrategy()
+  fetchStrategy: FetchStrategy = QueueSizeBasedFetchStrategy(),
+  metricLabels: Set[MetricLabel] = Set.empty,
+  runloopMetricsSchedule: Schedule[Any, Unit, Long] = Schedule.fixed(500.millis)
 ) {
 
   /**
@@ -278,6 +281,29 @@ final case class ConsumerSettings(
    */
   def withFetchStrategy(fetchStrategy: FetchStrategy): ConsumerSettings =
     copy(fetchStrategy = fetchStrategy)
+
+  /**
+   * @param metricLabels
+   *   The labels given to all metrics collected by zio-kafka. By default no labels are set.
+   *
+   * For applications with multiple consumers it is recommended to set some metric labels. For example, if one is used,
+   * the consumer group id could be used as a label:
+   *
+   * {{{
+   *   consumerSettings.withMetricLabels(Set(MetricLabel("group-id", groupId)))
+   * }}}
+   */
+  def withMetricsLabels(metricLabels: Set[MetricLabel]): ConsumerSettings =
+    copy(metricLabels = metricLabels)
+
+  /**
+   * @param runloopMetricsSchedule
+   *   The schedule at which the runloop metrics are measured. Example runloop metrics are queue sizes and number of
+   *   outstanding commits. The default is to measure every 500ms.
+   */
+  def withRunloopMetricsSchedule(runloopMetricsSchedule: Schedule[Any, Unit, Long]): ConsumerSettings =
+    copy(runloopMetricsSchedule = runloopMetricsSchedule)
+
 }
 
 object ConsumerSettings {