Remove lost partitions from assigned streams

Fixes #1288
zio · Oct 29, 2024 · b4f95d3 · b4f95d3
1 parent 4cd5733
commit b4f95d3
Showing 1 changed file with 4 additions and 3 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
@@ -236,7 +236,7 @@ private[consumer] final class Runloop private (
           state          <- currentStateRef.get
           lostStreams = state.assignedStreams.filter(control => lostTps.contains(control.tp))
           _ <- ZIO.foreachDiscard(lostStreams)(_.lost)
-          _ <- lastRebalanceEvent.set(rebalanceEvent.onLost(lostTps))
+          _ <- lastRebalanceEvent.set(rebalanceEvent.onLost(lostTps, lostStreams))
           _ <- ZIO.logTrace(s"onLost done")
         } yield ()
     )
@@ -830,11 +830,12 @@ object Runloop {
         endedStreams = this.endedStreams ++ endedStreams
       )
 
-    def onLost(lost: Set[TopicPartition]): RebalanceEvent =
+    def onLost(lost: Set[TopicPartition], endedStreams: Chunk[PartitionStreamControl]): RebalanceEvent =
       copy(
         wasInvoked = true,
         assignedTps = assignedTps -- lost,
-        lostTps = lostTps ++ lost
+        lostTps = lostTps ++ lost,
+        endedStreams = this.endedStreams ++ endedStreams
       )
   }