diff --git a/faust/events.py b/faust/events.py index e45dd4fe1..5c64c352f 100644 --- a/faust/events.py +++ b/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, diff --git a/faust/sensors/base.py b/faust/sensors/base.py index 46f679e67..6f1c350b8 100644 --- a/faust/sensors/base.py +++ b/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 diff --git a/faust/streams.py b/faust/streams.py index 9f4286a3f..9a775eb86 100644 --- a/faust/streams.py +++ b/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 diff --git a/faust/types/events.py b/faust/types/events.py index 9b1498609..33d4d807e 100644 --- a/faust/types/events.py +++ b/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__( diff --git a/tests/functional/test_streams.py b/tests/functional/test_streams.py index 851594155..4018ea608 100644 --- a/tests/functional/test_streams.py +++ b/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,27 @@ async def test_take(app): assert s.enable_acks is True +@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) diff --git a/tests/unit/test_streams.py b/tests/unit/test_streams.py index 4a6cab481..83149eb4a 100644 --- a/tests/unit/test_streams.py +++ b/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,