Fix #319: record event runtime for events acked via stream.take()

faust/events.py

@@ -2,7 +2,7 @@
 
 import typing
 from types import TracebackType
-from typing import Any, Awaitable, Optional, Type, Union, cast
+from typing import Any, Awaitable, Dict, Optional, Type, Union, cast
 
 from faust.types import (
     AppT,
@@ -128,6 +128,11 @@ def __init__(
             self.headers = {}
 
         self.acked: bool = False
+        #: Per-sensor state captured by ``Sensor.on_stream_event_in`` and
+        #: handed back to ``on_stream_event_out`` when the event is acked.
+        #: Stored here so deferred acks (e.g. ``Stream.take``) can still
+        #: report the event runtime.  See issue #319.
+        self.sensor_state: Optional[Dict] = None
 
     async def send(
         self,

faust/sensors/base.py

@@ -182,10 +182,16 @@ def on_stream_event_in(
         self, tp: TP, offset: int, stream: StreamT, event: EventT
     ) -> Optional[Dict]:
         """Call when stream starts processing an event."""
-        return {
+        state = {
             sensor: sensor.on_stream_event_in(tp, offset, stream, event)
             for sensor in self._sensors
         }
+        # Remember the per-sensor state on the event itself so that a
+        # deferred ack (e.g. Stream.take(), which acks manually long after
+        # the main loop yielded) can still deliver it to on_stream_event_out
+        # and record the true event runtime.  See issue #319.
+        event.sensor_state = state
+        return state
 
     def on_stream_event_out(
         self, tp: TP, offset: int, stream: StreamT, event: EventT, state: Dict = None

faust/streams.py

@@ -1268,7 +1268,13 @@ async def ack(self, event: EventT) -> bool:
         message = event.message
         tp = message.tp
         offset = message.offset
-        self._on_stream_event_out(tp, offset, self, event)
+        # Forward the sensor state captured in on_stream_event_in so the
+        # event runtime is recorded even when acking is deferred (e.g. the
+        # buffered events produced by Stream.take()).  Clear it afterwards so
+        # a second ack of the same event cannot double-count.  See issue #319.
+        sensor_state = getattr(event, "sensor_state", None)
+        self._on_stream_event_out(tp, offset, self, event, sensor_state)
+        event.sensor_state = None
         if last_stream_to_ack:
             self._on_message_out(tp, offset, message)
         return last_stream_to_ack

faust/types/events.py

@@ -4,6 +4,7 @@
     Any,
     AsyncContextManager,
     Awaitable,
+    Dict,
     Generic,
     Mapping,
     Optional,
@@ -38,8 +39,9 @@ class EventT(Generic[T], AsyncContextManager):
     headers: Mapping
     message: Message
     acked: bool
+    sensor_state: Optional[Dict]
 
-    __slots__ = ("app", "key", "value", "headers", "message", "acked")
+    __slots__ = ("app", "key", "value", "headers", "message", "acked", "sensor_state")
 
     @abc.abstractmethod
     def __init__(

tests/functional/test_streams.py

@@ -9,6 +9,7 @@
 
 import faust
 from faust.exceptions import ImproperlyConfigured
+from faust.sensors import Monitor
 from faust.streams import maybe_forward
 from tests.helpers import AsyncMock
 
@@ -777,6 +778,30 @@ async def test_take(app):
         assert s.enable_acks is True
 
 
+@pytest.mark.skipif(
+    platform.python_implementation() == "PyPy", reason="Not yet supported on PyPy"
+)
+@pytest.mark.asyncio
+async def test_take__records_event_runtime(app):
+    # Regression test for #319: events buffered by take() are acked manually
+    # after the main loop has already yielded, so on_stream_event_out used to
+    # be called without the sensor state and no event runtime was recorded.
+    monitor = Monitor()
+    app.sensors.add(monitor)
+    async with new_stream(app) as s:
+        await s.channel.send(value=1)
+        async for value in s.take(1, within=1):
+            assert value == [1]
+            break
+        # let take()'s finally ack the buffered event
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+
+    # on_stream_event_out ran with the captured sensor state, so a runtime
+    # was appended (would be empty before the fix).
+    assert monitor.events_runtime
+
+
 @pytest.mark.asyncio
 async def test_take__10(app, loop):
     s = new_stream(app)

tests/unit/test_streams.py

@@ -210,15 +210,21 @@ async def agent(stream):
     async def test_ack(self, *, stream):
         event = Mock()
         event.ack.return_value = True
+        event.sensor_state = {"sensor": "state"}
         stream._on_stream_event_out = Mock()
         stream._on_message_out = Mock()
         assert await stream.ack(event)
+        # the sensor state captured in on_stream_event_in is forwarded so
+        # deferred acks still record the event runtime (issue #319)
         stream._on_stream_event_out.assert_called_once_with(
             event.message.tp,
             event.message.offset,
             stream,
             event,
+            {"sensor": "state"},
         )
+        # ...and cleared afterwards so a second ack can't double-count
+        assert event.sensor_state is None
         stream._on_message_out.assert_called_once_with(
             event.message.tp,
             event.message.offset,