Refined implementation of #1022 (comment)

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

@@ -2,12 +2,11 @@ package zio.kafka.consumer
 
 import org.apache.kafka.clients.consumer.{ ConsumerGroupMetadata, ConsumerRecord }
 import org.apache.kafka.common.TopicPartition
+import zio.RIO
 import zio.kafka.serde.Deserializer
-import zio.{ RIO, Task }
 
 final case class CommittableRecord[K, V](
   record: ConsumerRecord[K, V],
-  private val commitHandle: Map[TopicPartition, Long] => Task[Unit],
   private val consumerGroupMetadata: Option[ConsumerGroupMetadata]
 ) {
   def deserializeWith[R, K1, V1](
@@ -39,26 +38,15 @@ final case class CommittableRecord[K, V](
   def timestamp: Long = record.timestamp()
 
   lazy val topicPartition: TopicPartition = new TopicPartition(record.topic(), record.partition())
-
-  def offset: Offset =
-    OffsetImpl(
-      topic = record.topic(),
-      partition = record.partition(),
-      offset = record.offset(),
-      commitHandle = commitHandle,
-      consumerGroupMetadata = consumerGroupMetadata
-    )
 }
 
 object CommittableRecord {
   def apply[K, V](
     record: ConsumerRecord[K, V],
-    commitHandle: Map[TopicPartition, Long] => Task[Unit],
     consumerGroupMetadata: Option[ConsumerGroupMetadata]
   ): CommittableRecord[K, V] =
     new CommittableRecord(
       record = record,
-      commitHandle = commitHandle,
       consumerGroupMetadata = consumerGroupMetadata
     )
 }

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

@@ -2,10 +2,16 @@ package zio.kafka.consumer
 
 import org.apache.kafka.common.TopicPartition
 import zio.kafka.utils.PendingCommit
-import zio.{Chunk, Promise, Ref, Schedule, Scope, Task, UIO, URIO, ZIO}
+import zio.{ Chunk, Promise, Ref, Schedule, Scope, Task, UIO, URIO, ZIO }
 
 trait Committer {
   def commit(records: Chunk[CommittableRecord[_, _]]): UIO[PendingCommit[Throwable, Unit]]
+
+  final def commitAndAwait(records: Chunk[CommittableRecord[_, _]]): Task[Unit] =
+    commit(records).flatMap(_.awaitCommit)
+
+  final def commitAndForget(records: Chunk[CommittableRecord[_, _]]): UIO[Unit] =
+    commit(records).unit
 }
 
 //noinspection ConvertExpressionToSAM

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

@@ -1,17 +1,12 @@
 package zio.kafka.consumer
 
-import org.apache.kafka.clients.consumer.{
-  Consumer => JConsumer,
-  ConsumerRecord,
-  OffsetAndMetadata,
-  OffsetAndTimestamp
-}
+import org.apache.kafka.clients.consumer.{ConsumerRecord, OffsetAndMetadata, OffsetAndTimestamp, Consumer => JConsumer}
 import org.apache.kafka.common._
 import zio._
 import zio.kafka.consumer.diagnostics.DiagnosticEvent.Finalization
 import zio.kafka.consumer.diagnostics.Diagnostics
-import zio.kafka.consumer.internal.{ ConsumerAccess, RunloopAccess }
-import zio.kafka.serde.{ Deserializer, Serde }
+import zio.kafka.consumer.internal.{ConsumerAccess, RunloopAccess}
+import zio.kafka.serde.{Deserializer, Serde}
 import zio.kafka.utils.SslHelper
 import zio.stream._
 
@@ -169,9 +164,6 @@ object Consumer {
   case object RunloopTimeout extends RuntimeException("Timeout in Runloop") with NoStackTrace
   case object CommitTimeout  extends RuntimeException("Commit timeout") with NoStackTrace
 
-  val offsetBatches: ZSink[Any, Nothing, Offset, Nothing, OffsetBatch] =
-    ZSink.foldLeft[Offset, OffsetBatch](OffsetBatch.empty)(_ add _)
-
   def live: RLayer[ConsumerSettings & Diagnostics, Consumer] =
     ZLayer.scoped {
       for {
@@ -232,6 +224,9 @@ object Consumer {
   ): RIO[Consumer, Map[TopicPartition, Option[OffsetAndMetadata]]] =
     ZIO.serviceWithZIO(_.committed(partitions, timeout))
 
+  def commitAccumBatch[R](commitschedule: Schedule[R, Any, Any]): URIO[R & Consumer, Committer] =
+    ZIO.serviceWithZIO[Consumer](_.commitAccumBatch(commitschedule))
+
   /**
    * Accessor method
    */
@@ -423,7 +418,6 @@ private[consumer] final class ConsumerLive private[consumer] (
   consumer: ConsumerAccess,
   runloopAccess: RunloopAccess
 ) extends Consumer {
-  import Consumer._
 
   override def assignment: Task[Set[TopicPartition]] =
     consumer.withConsumer(_.assignment().asScala.toSet)
@@ -525,15 +519,17 @@ private[consumer] final class ConsumerLive private[consumer] (
     f: ConsumerRecord[K, V] => URIO[R1, Unit]
   ): ZIO[R & R1, Throwable, Unit] =
     for {
-      r <- ZIO.environment[R & R1]
+      r         <- ZIO.environment[R & R1]
+      committer <- commitAccumBatch(commitRetryPolicy)
       _ <- partitionedStream(subscription, keyDeserializer, valueDeserializer)
              .flatMapPar(Int.MaxValue) { case (_, partitionStream) =>
-               partitionStream.mapChunksZIO(_.mapZIO((c: CommittableRecord[K, V]) => f(c.record).as(c.offset)))
+               partitionStream.mapChunksZIO(records =>
+                 records.mapZIO((r: CommittableRecord[K, V]) => f(r.record)).as(records)
+               )
              }
-             .provideEnvironment(r)
-             .aggregateAsync(offsetBatches)
-             .mapZIO(_.commitOrRetry(commitRetryPolicy))
+             .mapChunksZIO(committer.commitAndAwait(_).as(Chunk.empty))
              .runDrain
+             .provideEnvironment(r)
     } yield ()
 
   override def offsetsForTimes(