Make clearing rebalance event more obvious

zio · Oct 28, 2023 · 1ad29ef · 1ad29ef
1 parent 9f4b41b
commit 1ad29ef
Showing 1 changed file with 1 addition and 2 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
@@ -294,7 +294,6 @@ private[consumer] final class Runloop private (
                s" resuming ${partitionsToFetch} partitions"
            )
       _ <- currentStateRef.set(state)
-      _ <- lastRebalanceEvent.set(RebalanceEvent.None)
       pollResult <-
         consumer.runloopAccess { c =>
           ZIO.suspend {
@@ -317,7 +316,7 @@ private[consumer] final class Runloop private (
                 tpWithoutData = requestedPartitions -- providedTps
               )
             } *>
-              lastRebalanceEvent.get.flatMap {
+              lastRebalanceEvent.getAndSet(RebalanceEvent.None).flatMap {
                 case RebalanceEvent(false, _, _, _, _) =>
                   // The fast track, rebalance listener was not invoked:
                   //   no assignment changes, only new records.