diff --git a/faust/transport/consumer.py b/faust/transport/consumer.py index ab3554416..01e699983 100644 --- a/faust/transport/consumer.py +++ b/faust/transport/consumer.py @@ -761,7 +761,13 @@ async def getmany(self, timeout: float) -> AsyncIterator[Tuple[TP, Message]]: or tp in self._buffered_partitions ): highwater_mark = self.highwater(tp) - self.app.monitor.track_tp_end_offset(tp, highwater_mark) + # highwater() can return None during a rebalance before + # the end offset is known. Tracking it would crash the + # metric sensors that do float(offset) (Prometheus, + # Datadog, StatsD) and take down _drain_messages. Skip + # until the highwater is known. See issue #214. + if highwater_mark is not None: + self.app.monitor.track_tp_end_offset(tp, highwater_mark) # convert timestamp to seconds from int milliseconds. yield tp, to_message(tp, record) else: diff --git a/tests/unit/transport/test_consumer.py b/tests/unit/transport/test_consumer.py index 8a5986e58..086e32f55 100644 --- a/tests/unit/transport/test_consumer.py +++ b/tests/unit/transport/test_consumer.py @@ -509,6 +509,37 @@ def to_message(tp, record): (TP2, "G"), ] + @pytest.mark.asyncio + async def test_getmany__highwater_none_not_tracked(self, *, consumer): + # Regression test for #214: highwater() can return None during a + # rebalance. Passing it to track_tp_end_offset crashes metric + # sensors that do float(offset), so it must be skipped. + def to_message(tp, record): + return record + + consumer._to_message = to_message + consumer.highwater = Mock(name="highwater", return_value=None) + consumer.app.monitor = Mock(name="monitor") + self._setup_records(consumer, active_partitions={TP1}, records={TP1: ["A"]}) + consumer.flow_active = True + + assert [a async for a in consumer.getmany(1.0)] == [(TP1, "A")] + consumer.app.monitor.track_tp_end_offset.assert_not_called() + + @pytest.mark.asyncio + async def test_getmany__highwater_tracked(self, *, consumer): + def to_message(tp, record): + return record + + consumer._to_message = to_message + consumer.highwater = Mock(name="highwater", return_value=42) + consumer.app.monitor = Mock(name="monitor") + self._setup_records(consumer, active_partitions={TP1}, records={TP1: ["A"]}) + consumer.flow_active = True + + assert [a async for a in consumer.getmany(1.0)] == [(TP1, "A")] + consumer.app.monitor.track_tp_end_offset.assert_called_once_with(TP1, 42) + @pytest.mark.asyncio async def test_getmany_buffered(self, *, consumer): def to_message(tp, record):