Discard fetched messages on rebalance

faust/streams.py

@@ -974,6 +974,8 @@ async def _py_aiter(self) -> AsyncIterator[T_co]:
                         tp = message.tp
                         offset = message.offset
 
+                        if message.generation_id != self.app.consumer_generation_id:
+                            continue
                         if topic in acking_topics and not message.tracked:
                             message.tracked = True
                             # This inlines Consumer.track_message(message)

faust/transport/consumer.py

@@ -435,6 +435,7 @@ class Consumer(Service, ConsumerT):
 
     flow_active: bool = True
     can_resume_flow: Event
+    suspend_flow: Event
 
     def __init__(
         self,
@@ -475,6 +476,7 @@ def __init__(
         self._end_offset_monitor_interval = self.commit_interval * 2
         self.randomly_assigned_topics = set()
         self.can_resume_flow = Event()
+        self.suspend_flow = Event()
         self._reset_state()
         super().__init__(loop=loop or self.transport.loop, **kwargs)
         self.transactions = self.transport.create_transaction_manager(
@@ -497,6 +499,7 @@ def _reset_state(self) -> None:
         self._paused_partitions = set()
         self._buffered_partitions = set()
         self.can_resume_flow.clear()
+        self.suspend_flow.clear()
         self.flow_active = True
         self._time_start = monotonic()
 
@@ -573,11 +576,13 @@ def stop_flow(self) -> None:
         """Block consumer from processing any more messages."""
         self.flow_active = False
         self.can_resume_flow.clear()
+        self.suspend_flow.set()
 
     def resume_flow(self) -> None:
         """Allow consumer to process messages."""
         self.flow_active = True
         self.can_resume_flow.set()
+        self.suspend_flow.clear()
 
     def pause_partitions(self, tps: Iterable[TP]) -> None:
         """Pause fetching from partitions."""
@@ -1120,7 +1125,9 @@ async def _drain_messages(self, fetcher: ServiceT) -> None:  # pragma: no cover
                                 if self._n_acked >= commit_every:
                                     self._n_acked = 0
                                     await self.commit()
-                            await callback(message)
+                            await self.wait_first(
+                                callback(message), self.suspend_flow.wait()
+                            )
                             set_read_offset(tp, offset)
                         else:
                             self.log.dev(

faust/transport/drivers/aiokafka.py

@@ -267,6 +267,7 @@ def _to_message(self, tp: TP, record: Any) -> ConsumerMessage:
             record.serialized_key_size,
             record.serialized_value_size,
             tp,
+            generation_id=self.app.consumer_generation_id,
         )
 
     async def on_stop(self) -> None:

faust/types/tuples.py

@@ -133,6 +133,7 @@ class Message:
         "tracked",
         "span",
         "__weakref__",
+        "generation_id",
     )
 
     use_tracking: bool = False
@@ -154,6 +155,7 @@ def __init__(
         time_in: float = None,
         time_out: float = None,
         time_total: float = None,
+        generation_id: int = None,
     ) -> None:
         self.topic: str = topic
         self.partition: int = partition
@@ -183,6 +185,12 @@ def __init__(
         #: still processing.
         self.time_total: Optional[float] = time_total
 
+        # In some edge cases a message can slip through to the stream from before a
+        # rebalance occured if it gets stuck in the conductor or somewhere else. We
+        # track the generation_id when the message is fetched so we can discard if
+        # needed.
+        self.generation_id: Optional[int] = generation_id
+
     def ack(self, consumer: _ConsumerT, n: int = 1) -> bool:
         if not self.acked:
             # if no more references, mark offset as safe-to-commit in