Add take_events()

Author: nguyentritai2906

faust/streams.py

@@ -206,7 +206,8 @@ def _iter_ll(self, dir_: _LinkedListDirection) -> Iterator[StreamT]:
         seen: Set[StreamT] = set()
         while node:
             if node in seen:
-                raise RuntimeError(f"Loop in Stream.{dir_.attr}: Call support!")
+                raise RuntimeError(
+                    f"Loop in Stream.{dir_.attr}: Call support!")
             seen.add(node)
             yield node
             node = dir_.getter(node)
@@ -270,9 +271,7 @@ def _chain(self, **kwargs: Any) -> StreamT:
 
     def noack(self) -> "StreamT":
         """Create new stream where acks are manual."""
-        self._next = new_stream = self.clone(
-            enable_acks=False,
-        )
+        self._next = new_stream = self.clone(enable_acks=False, )
         return new_stream
 
     async def items(self) -> AsyncIterator[Tuple[K, T_co]]:
@@ -299,7 +298,8 @@ async def events(self) -> AsyncIterable[EventT]:
             if self.current_event is not None:
                 yield self.current_event
 
-    async def take(self, max_: int, within: Seconds) -> AsyncIterable[Sequence[T_co]]:
+    async def take(self, max_: int,
+                   within: Seconds) -> AsyncIterable[Sequence[T_co]]:
         """Buffer n values at a time and yield a list of buffered values.
 
         Arguments:
@@ -340,7 +340,8 @@ async def add_to_buffer(value: T) -> T:
                 buffer_add(cast(T_co, value))
                 event = self.current_event
                 if event is None:
-                    raise RuntimeError("Take buffer found current_event is None")
+                    raise RuntimeError(
+                        "Take buffer found current_event is None")
                 event_add(event)
                 if buffer_size() >= max_:
                     # signal that the buffer is full and should be emptied.
@@ -391,9 +392,102 @@ async def add_to_buffer(value: T) -> T:
             self.enable_acks = stream_enable_acks
             self._processors.remove(add_to_buffer)
 
+    async def take_events(self, max_: int,
+                          within: Seconds) -> AsyncIterable[Sequence[EventT]]:
+        """Buffer n events at a time and yield a list of buffered events.
+        Arguments:
+            max_: Max number of messages to receive. When more than this
+                number of messages are received within the specified number of
+                seconds then we flush the buffer immediately.
+            within: Timeout for when we give up waiting for another value,
+                and process the values we have.
+                Warning: If there's no timeout (i.e. `timeout=None`),
+                the agent is likely to stall and block buffered events for an
+                unreasonable length of time(!).
+        """
+        buffer: List[T_co] = []
+        events: List[EventT] = []
+        buffer_add = buffer.append
+        event_add = events.append
+        buffer_size = buffer.__len__
+        buffer_full = asyncio.Event()
+        buffer_consumed = asyncio.Event()
+        timeout = want_seconds(within) if within else None
+        stream_enable_acks: bool = self.enable_acks
+
+        buffer_consuming: Optional[asyncio.Future] = None
+
+        channel_it = aiter(self.channel)
+
+        # We add this processor to populate the buffer, and the stream
+        # is passively consumed in the background (enable_passive below).
+        async def add_to_buffer(value: T) -> T:
+            try:
+                # buffer_consuming is set when consuming buffer after timeout.
+                nonlocal buffer_consuming
+                if buffer_consuming is not None:
+                    try:
+                        await buffer_consuming
+                    finally:
+                        buffer_consuming = None
+                buffer_add(cast(T_co, value))
+                event = self.current_event
+                if event is None:
+                    raise RuntimeError(
+                        'Take buffer found current_event is None')
+                event_add(event)
+                if buffer_size() >= max_:
+                    # signal that the buffer is full and should be emptied.
+                    buffer_full.set()
+                    # strict wait for buffer to be consumed after buffer full.
+                    # If max is 1000, we are not allowed to return 1001 values.
+                    buffer_consumed.clear()
+                    await self.wait(buffer_consumed)
+            except CancelledError:  # pragma: no cover
+                raise
+            except Exception as exc:
+                self.log.exception('Error adding to take buffer: %r', exc)
+                await self.crash(exc)
+            return value
+
+        # Disable acks to ensure this method acks manually
+        # events only after they are consumed by the user
+        self.enable_acks = False
+
+        self.add_processor(add_to_buffer)
+        self._enable_passive(cast(ChannelT, channel_it))
+        try:
+            while not self.should_stop:
+                # wait until buffer full, or timeout
+                await self.wait_for_stopped(buffer_full, timeout=timeout)
+                if buffer:
+                    # make sure background thread does not add new items to
+                    # buffer while we read.
+                    buffer_consuming = self.loop.create_future()
+                    try:
+                        yield list(events)
+                    finally:
+                        buffer.clear()
+                        for event in events:
+                            await self.ack(event)
+                        events.clear()
+                        # allow writing to buffer again
+                        notify(buffer_consuming)
+                        buffer_full.clear()
+                        buffer_consumed.set()
+                else:  # pragma: no cover
+                    pass
+            else:  # pragma: no cover
+                pass
+
+        finally:
+            # Restore last behaviour of "enable_acks"
+            self.enable_acks = stream_enable_acks
+            self._processors.remove(add_to_buffer)
+
     async def take_with_timestamp(
-        self, max_: int, within: Seconds, timestamp_field_name: str
-    ) -> AsyncIterable[Sequence[T_co]]:
+            self, max_: int, within: Seconds,
+            timestamp_field_name: str) -> AsyncIterable[Sequence[T_co]]:
         """Buffer n values at a time and yield a list of buffered values with the
            timestamp when the message was added to kafka.
 
@@ -439,7 +533,8 @@ async def add_to_buffer(value: T) -> T:
                     value[timestamp_field_name] = event.message.timestamp
                 buffer_add(value)
                 if event is None:
-                    raise RuntimeError("Take buffer found current_event is None")
+                    raise RuntimeError(
+                        "Take buffer found current_event is None")
                 event_add(event)
                 if buffer_size() >= max_:
                     # signal that the buffer is full and should be emptied.
@@ -498,9 +593,8 @@ def enumerate(self, start: int = 0) -> AsyncIterable[Tuple[int, T_co]]:
         """
         return aenumerate(self, start)
 
-    async def noack_take(
-        self, max_: int, within: Seconds
-    ) -> AsyncIterable[Sequence[T_co]]:
+    async def noack_take(self, max_: int,
+                         within: Seconds) -> AsyncIterable[Sequence[T_co]]:
         """
          Buffer n values at a time and yield a list of buffered values.
         :param max_: Max number of messages to receive. When more than this
@@ -543,7 +637,8 @@ async def add_to_buffer(value: T) -> T:
                 event = self.current_event
                 buffer_add(cast(T_co, event))
                 if event is None:
-                    raise RuntimeError("Take buffer found current_event is None")
+                    raise RuntimeError(
+                        "Take buffer found current_event is None")
 
                 event_add(event)
                 if buffer_size() >= max_:
@@ -628,8 +723,7 @@ async def mytask(stream):
             return self
         if self.concurrency_index is not None:
             raise ImproperlyConfigured(
-                "Agent with concurrency>1 cannot use stream.through!"
-            )
+                "Agent with concurrency>1 cannot use stream.through!")
         # ridiculous mypy
         if isinstance(channel, str):
             channelchannel = cast(ChannelT, self.derive_topic(channel))
@@ -638,7 +732,8 @@ async def mytask(stream):
 
         channel_it = aiter(channelchannel)
         if self._next is not None:
-            raise ImproperlyConfigured("Stream is already using group_by/through")
+            raise ImproperlyConfigured(
+                "Stream is already using group_by/through")
         through = self._chain(channel=channel_it)
 
         async def forward(value: T) -> T:
@@ -649,12 +744,17 @@ async def forward(value: T) -> T:
         self._enable_passive(cast(ChannelT, channel_it), declare=True)
         return through
 
-    def _enable_passive(self, channel: ChannelT, *, declare: bool = False) -> None:
+    def _enable_passive(self,
+                        channel: ChannelT,
+                        *,
+                        declare: bool = False) -> None:
         if not self._passive:
             self._passive = True
             self.add_future(self._passive_drainer(channel, declare))
 
-    async def _passive_drainer(self, channel: ChannelT, declare: bool = False) -> None:
+    async def _passive_drainer(self,
+                               channel: ChannelT,
+                               declare: bool = False) -> None:
         try:
             if declare:
                 await channel.maybe_declare()
@@ -760,13 +860,13 @@ def get_key(withdrawal):
         channel: ChannelT
         if self.concurrency_index is not None:
             raise ImproperlyConfigured(
-                "Agent with concurrency>1 cannot use stream.group_by!"
-            )
+                "Agent with concurrency>1 cannot use stream.group_by!")
         if not name:
             if isinstance(key, FieldDescriptorT):
                 name = key.ident
             else:
-                raise TypeError("group_by with callback must set name=topic_suffix")
+                raise TypeError(
+                    "group_by with callback must set name=topic_suffix")
         if topic is not None:
             channel = topic
         else:
@@ -776,9 +876,10 @@ def get_key(withdrawal):
                 prefix = self.prefix + "-"
             suffix = f"-{name}-repartition"
             p = partitions if partitions else self.app.conf.topic_partitions
-            channel = cast(ChannelT, self.channel).derive(
-                prefix=prefix, suffix=suffix, partitions=p, internal=True
-            )
+            channel = cast(ChannelT, self.channel).derive(prefix=prefix,
+                                                          suffix=suffix,
+                                                          partitions=p,
+                                                          internal=True)
         format_key = self._format_key
 
         channel_it = aiter(channel)
@@ -789,7 +890,8 @@ def get_key(withdrawal):
         async def repartition(value: T) -> T:
             event = self.current_event
             if event is None:
-                raise RuntimeError("Cannot repartition stream with non-topic channel")
+                raise RuntimeError(
+                    "Cannot repartition stream with non-topic channel")
             new_key = await format_key(key, value)
             await event.forward(channel, key=new_key)
             return value
@@ -1070,10 +1172,9 @@ async def _py_aiter(self) -> AsyncIterator[T_co]:
                         tp = message.tp
                         offset = message.offset
 
-                        if (
-                            not self.app.flow_control.is_active()
-                            or message.generation_id != self.app.consumer_generation_id
-                        ):
+                        if (not self.app.flow_control.is_active()
+                                or message.generation_id !=
+                                self.app.consumer_generation_id):
                             value = skipped_value
                             self.log.dev(
                                 "Skipping message %r with generation_id %r because "
@@ -1092,7 +1193,8 @@ async def _py_aiter(self) -> AsyncIterator[T_co]:
                             # XXX ugh this should be in the consumer somehow
 
                         # call Sensors
-                        sensor_state = on_stream_event_in(tp, offset, self, event)
+                        sensor_state = on_stream_event_in(
+                            tp, offset, self, event)
 
                         # set task-local current_event
                         _current_event_contextvar.set(create_ref(event))
@@ -1125,13 +1227,15 @@ async def _py_aiter(self) -> AsyncIterator[T_co]:
                     self.current_event = None
                     # We want to ack the filtered out message
                     # otherwise the lag would increase
-                    if event is not None and (do_ack or value is skipped_value):
+                    if event is not None and (do_ack
+                                              or value is skipped_value):
                         # This inlines self.ack
                         last_stream_to_ack = event.ack()
                         message = event.message
                         tp = event.message.tp
                         offset = event.message.offset
-                        on_stream_event_out(tp, offset, self, event, sensor_state)
+                        on_stream_event_out(tp, offset, self, event,
+                                            sensor_state)
                         if last_stream_to_ack:
                             on_message_out(tp, offset, message)