async

zio · Jun 11, 2023 · 7a18256 · 7a18256
1 parent de5c031
commit 7a18256
Showing 1 changed file with 31 additions and 9 deletions.
diff --git a/zio-kafka/src/main/scala/zio/kafka/consumer/internal/Runloop.scala b/zio-kafka/src/main/scala/zio/kafka/consumer/internal/Runloop.scala
@@ -12,6 +12,7 @@ import zio.kafka.consumer.internal.ConsumerAccess.ByteArrayKafkaConsumer
 import zio.kafka.consumer.internal.RunloopAccess.PartitionAssignment
 import zio.stream._
 
+import java.util
 import scala.jdk.CollectionConverters._
 
 //noinspection SimplifyWhenInspection
@@ -118,15 +119,36 @@ private[consumer] final class Runloop private (
     val offsets = cmd.offsets.map { case (tp, offset) => tp -> new OffsetAndMetadata(offset + 1) }
 
     consumer
-      .runloopAccess(c => ZIO.attempt(c.commitSync(offsets.asJava)))
-      .foldZIO(
-        {
-          case _: RebalanceInProgressException =>
-            ZIO.logDebug(s"Rebalance in progress, retrying commit for offsets $offsets") *>
-              commandQueue.offer(cmd).unit
-          case err => cmd.fail(err) <* diagnostics.emit(DiagnosticEvent.Commit.Failure(offsets, err))
-        },
-        _ => cmd.succeed <* diagnostics.emit(DiagnosticEvent.Commit.Success(offsets))
+      .runloopAccess(c =>
+        for {
+          fiberId <- ZIO.fiberId
+          _       <- ZIO.logDebug(s"fiberId: $fiberId")
+          _ <- ZIO.async(
+                 register = (cb: UIO[Unit] => Unit) =>
+                   c.commitAsync(
+                     offsets.asJava,
+                     (_: util.Map[TopicPartition, OffsetAndMetadata], e: Exception) =>
+                       if (e == null)
+                         cb(
+                           ZIO.logDebug(s"Done for: $offsets") *> cmd.succeed <*
+                             diagnostics.emit(DiagnosticEvent.Commit.Success(offsets))
+                         )
+                       else
+                         cb(
+                           ZIO.logDebugCause(s"Failed for: $offsets", Cause.fail(e)) *>
+                             (e match {
+                               case _: RebalanceInProgressException =>
+                                 ZIO.logDebug(s"Rebalance in progress, retrying commit for offsets $offsets") *>
+                                   commandQueue.offer(cmd).unit
+                               case err =>
+                                 ZIO.logDebug(s"TATA") *> cmd.fail(err) <*
+                                   diagnostics.emit(DiagnosticEvent.Commit.Failure(offsets, err))
+                             })
+                         )
+                   ),
+                 blockingOn = fiberId
+               )
+        } yield ()
       )
   }