Extract Committer and RebalanceCoordinator classes from Runloop + unit tests

zio-kafka-test/src/test/scala/zio/kafka/consumer/ConsumerSpec.scala

@@ -28,6 +28,7 @@ import zio.test.Assertion._
 import zio.test.TestAspect._
 import zio.test._
 
+import java.util.concurrent.ExecutionException
 import scala.reflect.ClassTag
 
 //noinspection SimplifyAssertInspection
@@ -769,7 +770,8 @@ object ConsumerSpec extends ZIOSpecDefaultSlf4j with KafkaRandom {
                               clientId = clientId,
                               groupId = Some(groupId),
                               `max.poll.records` = 1,
-                              rebalanceSafeCommits = rebalanceSafeCommits
+                              rebalanceSafeCommits = rebalanceSafeCommits,
+                              maxRebalanceDuration = 60.seconds
                             )
                 consumer <- Consumer.make(settings)
               } yield consumer
@@ -1376,9 +1378,13 @@ object ConsumerSpec extends ZIOSpecDefaultSlf4j with KafkaRandom {
       }: _*) @@ TestAspect.nonFlaky(2),
       test("running streams don't stall after a poll timeout") {
         for {
-          topic      <- randomTopic
-          clientId   <- randomClient
-          _          <- ZIO.fromTry(EmbeddedKafka.createCustomTopic(topic))
+          topic    <- randomTopic
+          clientId <- randomClient
+          _ <- ZIO.attempt(EmbeddedKafka.createCustomTopic(topic)).flatMap(ZIO.fromTry(_)).retryN(3).catchSome {
+                 case e: ExecutionException
+                     if e.getCause.isInstanceOf[org.apache.kafka.common.errors.TopicExistsException] =>
+                   ZIO.unit
+               }
           settings   <- consumerSettings(clientId)
           consumer   <- Consumer.make(settings.withPollTimeout(50.millis))
           recordsOut <- Queue.unbounded[Unit]

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

@@ -1,11 +1,12 @@
 package zio.kafka.consumer.internal
 
+import org.apache.kafka.clients.consumer.OffsetAndMetadata
 import org.apache.kafka.common.TopicPartition
 import zio._
-import org.apache.kafka.clients.consumer.OffsetAndMetadata
+import zio.kafka.consumer.internal.Committer.{ Commit, CommitOffsets }
 import zio.test._
 
-object RunloopCommitOffsetsSpec extends ZIOSpecDefault {
+object CommitOffsetsSpec extends ZIOSpecDefault {
 
   private val tp10 = new TopicPartition("t1", 0)
   private val tp11 = new TopicPartition("t1", 1)
@@ -14,49 +15,49 @@ object RunloopCommitOffsetsSpec extends ZIOSpecDefault {
   private val tp22 = new TopicPartition("t2", 2)
 
   override def spec: Spec[TestEnvironment with Scope, Any] =
-    suite("Runloop.CommitOffsets spec")(
+    suite("CommitOffsets spec")(
       test("addCommits adds to empty CommitOffsets") {
-        val s1        = Runloop.CommitOffsets(Map.empty)
+        val s1        = 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 -> 4L))
+        val s1        = 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 s1        = 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 s1        = 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 s1        = 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 s1 = CommitOffsets(Map(tp10 -> 10L, tp20 -> 200L, tp21 -> 210L, tp22 -> 220L))
         val (inc, s2) = s1.addCommits(
           Chunk(
             makeCommit(Map(tp11 -> 11, tp20 -> 199L, tp21 -> 211L, tp22 -> 219L)),
@@ -69,35 +70,35 @@ object RunloopCommitOffsetsSpec extends ZIOSpecDefault {
         )
       },
       test("keepPartitions removes some partitions") {
-        val s1 = Runloop.CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
+        val s1 = CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
         val s2 = s1.keepPartitions(Set(tp10))
         assertTrue(s2.offsets == Map(tp10 -> 10L))
       },
       test("does not 'contain' offset when tp is not present") {
-        val s1     = Runloop.CommitOffsets(Map(tp10 -> 10L))
+        val s1     = CommitOffsets(Map(tp10 -> 10L))
         val result = s1.contains(tp20, 10)
         assertTrue(!result)
       },
       test("does not 'contain' a higher offset") {
-        val s1     = Runloop.CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
+        val s1     = CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
         val result = s1.contains(tp10, 11)
         assertTrue(!result)
       },
       test("does 'contain' equal offset") {
-        val s1     = Runloop.CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
+        val s1     = CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
         val result = s1.contains(tp10, 10)
         assertTrue(result)
       },
       test("does 'contain' lower offset") {
-        val s1     = Runloop.CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
+        val s1     = CommitOffsets(Map(tp10 -> 10L, tp20 -> 20L))
         val result = s1.contains(tp20, 19)
         assertTrue(result)
       }
     )
 
-  private def makeCommit(offsets: Map[TopicPartition, Long]): Runloop.Commit = {
+  private def makeCommit(offsets: Map[TopicPartition, Long]): 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(0L, o, p)
+    Commit(0L, o, p)
   }
 }

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

@@ -0,0 +1,207 @@
+package zio.kafka.consumer.internal
+import org.apache.kafka.clients.consumer.{ OffsetAndMetadata, OffsetCommitCallback }
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.errors.RebalanceInProgressException
+import zio.kafka.consumer.Consumer.CommitTimeout
+import zio.kafka.consumer.ConsumerSettings
+import zio.kafka.consumer.diagnostics.{ DiagnosticEvent, Diagnostics }
+import zio.kafka.consumer.internal.Committer.{ Commit, CommitOffsets }
+import zio.kafka.consumer.internal.ConsumerAccess.ByteArrayKafkaConsumer
+import zio.{ durationLong, Chunk, Exit, Promise, Queue, Ref, Runtime, Scope, Task, UIO, Unsafe, ZIO }
+
+import java.lang.Math.max
+import java.util
+import java.util.{ Map => JavaMap }
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+private[consumer] final class Committer(
+  commitQueue: Queue[Commit],
+  settings: ConsumerSettings,
+  diagnostics: Diagnostics,
+  consumerMetrics: ConsumerMetrics,
+  onCommitAvailable: UIO[Unit],
+  committedOffsetsRef: Ref[CommitOffsets],
+  sameThreadRuntime: Runtime[Any],
+  pendingCommits: Ref[Chunk[Commit]] // TODO make Commit internal
+) {
+  private val commitTimeout = settings.commitTimeout
+
+  /** This is the implementation behind the user facing api `Offset.commit`. */
+  val commit: Map[TopicPartition, OffsetAndMetadata] => Task[Unit] =
+    offsets =>
+      for {
+        p <- Promise.make[Throwable, Unit]
+        startTime = java.lang.System.nanoTime()
+        _ <- commitQueue.offer(Commit(java.lang.System.nanoTime(), offsets, p))
+        _ <- onCommitAvailable
+        _ <- diagnostics.emit(DiagnosticEvent.Commit.Started(offsets))
+        _ <- p.await.timeoutFail(CommitTimeout)(commitTimeout)
+        endTime = java.lang.System.nanoTime()
+        latency = (endTime - startTime).nanoseconds
+        _ <- consumerMetrics.observeCommit(latency)
+      } yield ()
+
+  /**
+   * Takes commits from the queue, commits them and adds them to pending commits
+   *
+   * If the queue is empty, nothing happens, unless executeOnEmpty is true.
+   *
+   * @param consumer
+   *   Consumer with exclusive access
+   */
+  def handleNewCommits(consumer: ByteArrayKafkaConsumer, executeOnEmpty: Boolean = false): Task[Unit] = for {
+    commits <- commitQueue.takeAll
+    _       <- ZIO.logDebug(s"Processing ${commits.size} commits")
+    _ <- ZIO.unless(commits.isEmpty && !executeOnEmpty) {
+           val (offsets, callback, onFailure) = asyncCommitParameters(commits)
+           pendingCommits.update(_ ++ commits) *>
+             // We don't wait for the completion of the commit here, because it
+             // will only complete once we poll again.
+             ZIO
+               .attempt(consumer.commitAsync(offsets, callback))
+               .catchAll(onFailure)
+         }
+  } yield ()
+
+  /** Merge commits and prepare parameters for calling `consumer.commitAsync`. */
+  private def asyncCommitParameters(
+    commits: Chunk[Commit]
+  ): (JavaMap[TopicPartition, OffsetAndMetadata], OffsetCommitCallback, Throwable => UIO[Unit]) = {
+    val offsets = commits
+      .foldLeft(mutable.Map.empty[TopicPartition, OffsetAndMetadata]) { case (acc, commit) =>
+        commit.offsets.foreach { case (tp, offset) =>
+          acc += (tp -> acc
+            .get(tp)
+            .map(current => if (current.offset() > offset.offset()) current else offset)
+            .getOrElse(offset))
+        }
+        acc
+      }
+      .toMap
+    val offsetsWithMetaData = offsets.map { case (tp, offset) =>
+      tp -> new OffsetAndMetadata(offset.offset + 1, offset.leaderEpoch, offset.metadata)
+    }
+    val cont = (e: Exit[Throwable, Unit]) => ZIO.foreachDiscard(commits)(_.cont.done(e))
+    // We assume the commit is started immediately after returning from this method.
+    val startTime = java.lang.System.nanoTime()
+    val onSuccess = {
+      val endTime = java.lang.System.nanoTime()
+      val latency = (endTime - startTime).nanoseconds
+      for {
+        offsetIncrease <- committedOffsetsRef.modify(_.addCommits(commits))
+        _              <- consumerMetrics.observeAggregatedCommit(latency, offsetIncrease).when(commits.nonEmpty)
+        result         <- cont(Exit.unit)
+        _              <- diagnostics.emit(DiagnosticEvent.Commit.Success(offsetsWithMetaData))
+      } yield result
+    }
+    val onFailure: Throwable => UIO[Unit] = {
+      case _: RebalanceInProgressException =>
+        for {
+          _ <- ZIO.logDebug(s"Rebalance in progress, commit for offsets $offsets will be retried")
+          _ <- commitQueue.offerAll(commits)
+          _ <- onCommitAvailable //
+        } yield ()
+      case err: Throwable =>
+        cont(Exit.fail(err)) <* diagnostics.emit(DiagnosticEvent.Commit.Failure(offsetsWithMetaData, err))
+    }
+    val callback =
+      new OffsetCommitCallback {
+        override def onComplete(offsets: util.Map[TopicPartition, OffsetAndMetadata], exception: Exception): Unit =
+          Unsafe.unsafe { implicit u =>
+            sameThreadRuntime.unsafe.run {
+              if (exception eq null) onSuccess else onFailure(exception)
+            }
+              .getOrThrowFiberFailure()
+          }
+      }
+    (offsetsWithMetaData.asJava, callback, onFailure)
+  }
+
+  def pendingCommitCount: UIO[Int] = pendingCommits.get.map(_.size)
+
+  def queueSize: UIO[Int] = commitQueue.size
+
+  def updatePendingCommitsAfterPoll: UIO[Unit] =
+    pendingCommits.get.flatMap(ZIO.filter(_)(_.isPending)).flatMap(pendingCommits.set)
+
+  def pruneCommittedOffsets(assignedPartitions: Set[TopicPartition]): UIO[Unit] =
+    committedOffsetsRef.update(_.keepPartitions(assignedPartitions))
+
+  def getCommittedOffsets: UIO[CommitOffsets] = committedOffsetsRef.get
+
+  def getPendingCommits: UIO[CommitOffsets] =
+    pendingCommits.get.map(CommitOffsets.empty.addCommits(_)._2)
+}
+
+private[internal] object Committer {
+  def make(
+    settings: ConsumerSettings,
+    diagnostics: Diagnostics,
+    consumerMetrics: ConsumerMetrics,
+    onCommitAvailable: UIO[Unit],
+    sameThreadRuntime: Runtime[Any]
+  ): ZIO[Any with Scope, Nothing, Committer] = for {
+    pendingCommits      <- Ref.make(Chunk.empty[Commit])
+    commitQueue         <- ZIO.acquireRelease(Queue.unbounded[Commit])(_.shutdown)
+    committedOffsetsRef <- Ref.make(CommitOffsets.empty)
+  } yield new Committer(
+    commitQueue,
+    settings,
+    diagnostics,
+    consumerMetrics,
+    onCommitAvailable,
+    committedOffsetsRef,
+    sameThreadRuntime,
+    pendingCommits
+  )
+
+  private[internal] final case class Commit(
+    createdAt: NanoTime,
+    offsets: Map[TopicPartition, OffsetAndMetadata],
+    cont: Promise[Throwable, Unit]
+  ) {
+    @inline def isDone: UIO[Boolean]    = cont.isDone
+    @inline def isPending: UIO[Boolean] = isDone.negate
+  }
+
+  // package private for unit testing
+  private[internal] final case class CommitOffsets(offsets: Map[TopicPartition, Long]) {
+
+    /** Returns an estimate of the total offset increase, and a new `CommitOffsets` with the given offsets added. */
+    def addCommits(c: Chunk[Commit]): (Long, CommitOffsets) = {
+      val updatedOffsets = mutable.Map.empty[TopicPartition, Long]
+      updatedOffsets.sizeHint(offsets.size)
+      updatedOffsets ++= offsets
+      var offsetIncrease = 0L
+      c.foreach { commit =>
+        commit.offsets.foreach { case (tp, offsetAndMeta) =>
+          val offset = offsetAndMeta.offset()
+          val maxOffset = updatedOffsets.get(tp) match {
+            case Some(existingOffset) =>
+              offsetIncrease += max(0L, offset - existingOffset)
+              max(existingOffset, offset)
+            case None =>
+              // This partition was not committed to from this consumer yet. Therefore we do not know the offset
+              // increase. A good estimate would be the poll size for this consumer, another okayish estimate is 0.
+              // Lets go with the simplest for now: ```offsetIncrease += 0```
+              offset
+          }
+          updatedOffsets += tp -> maxOffset
+        }
+      }
+      (offsetIncrease, CommitOffsets(offsets = updatedOffsets.toMap))
+    }
+
+    def keepPartitions(tps: Set[TopicPartition]): CommitOffsets =
+      CommitOffsets(offsets.filter { case (tp, _) => tps.contains(tp) })
+
+    def contains(tp: TopicPartition, offset: Long): Boolean =
+      offsets.get(tp).exists(_ >= offset)
+  }
+
+  private[internal] object CommitOffsets {
+    val empty: CommitOffsets = CommitOffsets(Map.empty)
+  }
+
+}

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

@@ -16,7 +16,12 @@ private[internal] trait ConsumerMetrics {
   def observeCommit(latency: Duration): UIO[Unit]
   def observeAggregatedCommit(latency: Duration, commitSize: Long): UIO[Unit]
   def observeRebalance(currentlyAssignedCount: Int, assignedCount: Int, revokedCount: Int, lostCount: Int): UIO[Unit]
-  def observeRunloopMetrics(state: Runloop.State, commandQueueSize: Int, commitQueueSize: Int): UIO[Unit]
+  def observeRunloopMetrics(
+    state: Runloop.State,
+    commandQueueSize: Int,
+    commitQueueSize: Int,
+    pendingCommits: Int
+  ): UIO[Unit]
   def observePollAuthError(): UIO[Unit]
 }
 
@@ -330,14 +335,19 @@ private[internal] class ZioConsumerMetrics(metricLabels: Set[MetricLabel]) exten
       .contramap[Int](_.toDouble)
       .tagged(metricLabels)
 
-  override def observeRunloopMetrics(state: Runloop.State, commandQueueSize: Int, commitQueueSize: Int): UIO[Unit] =
+  override def observeRunloopMetrics(
+    state: Runloop.State,
+    commandQueueSize: Int,
+    commitQueueSize: Int,
+    pendingCommits: Int
+  ): UIO[Unit] =
     for {
       _          <- ZIO.foreachDiscard(state.assignedStreams)(_.outstandingPolls @@ queuePollsHistogram)
       queueSizes <- ZIO.foreach(state.assignedStreams)(_.queueSize)
       _          <- ZIO.foreachDiscard(queueSizes)(qs => queueSizeHistogram.update(qs))
       _          <- allQueueSizeHistogram.update(queueSizes.sum)
       _          <- pendingRequestsHistogram.update(state.pendingRequests.size)
-      _          <- pendingCommitsHistogram.update(state.pendingCommits.size)
+      _          <- pendingCommitsHistogram.update(pendingCommits)
       _          <- subscriptionStateGauge.update(state.subscriptionState)
       _          <- commandQueueSizeHistogram.update(commandQueueSize)
       _          <- commitQueueSizeHistogram.update(commitQueueSize)