Fix #214: skip end-offset metric when highwater is unknown (None)

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:

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):