async def commit(self,
topics: TPorTopicSet = None,
start_new_transaction: bool = True) -> bool:
"""Maybe commit the offset for all or specific topics.
Arguments:
topics: Set containing topics and/or TopicPartitions to commit.
"""
if self.app.client_only:
# client only cannot commit as consumer does not have group_id
return False
if await self.maybe_wait_for_commit_to_finish():
# original commit finished, return False as we did not commit
return False
self._commit_fut = asyncio.Future(loop=self.loop)
try:
return await self.force_commit(
topics,
start_new_transaction=start_new_transaction,
)
finally:
# set commit_fut to None so that next call will commit.
fut, self._commit_fut = self._commit_fut, None
# notify followers that the commit is done.
notify(fut)
async def _subscriber(self) -> None: # pragma: no cover
# the first time we start, we will wait two seconds
# to give agents a chance to start up and register their
# streams. This way we won't have N subscription requests at the
# start.
await self.sleep(2.0)
if not self.should_stop:
# tell the consumer to subscribe to the topics.
await self.app.consumer.subscribe(await self._update_indices())
notify(self._subscription_done)
# Now we wait for changes
ev = self._subscription_changed = asyncio.Event(loop=self.loop)
while not self.should_stop:
# Wait for something to add/remove topics from subscription.
await ev.wait()
if self.app.rebalancing:
# we do not want to perform a resubscribe if the application
# is rebalancing.
ev.clear()
else:
# The change could be in reaction to something like "all agents
# restarting", in that case it would be bad if we resubscribe
# over and over, so we wait for 45 seconds to make sure any
# further subscription requests will happen during the same
# rebalance.
await self.sleep(self._resubscribe_sleep_lock_seconds)
subscribed_topics = await self._update_indices()
await self.app.consumer.subscribe(subscribed_topics)
# clear the subscription_changed flag, so we can wait on it again.
ev.clear()
# wake-up anything waiting for the subscription to be done.
notify(self._subscription_done)
def produce(self, topic: str, key: bytes, value: bytes, partition: int,
on_delivery: Callable) -> None:
if self._producer is None:
raise RuntimeError('Producer not started')
if partition is not None:
self._producer.produce(
topic, key, value, partition, on_delivery=on_delivery,
)
else:
self._producer.produce(
topic, key, value, on_delivery=on_delivery,
)
notify(self._flush_soon)
def ack(self, message: Message) -> bool:
if not message.acked:
message.acked = True
tp = message.tp
offset = message.offset
if self.app.topics.acks_enabled_for(message.topic):
committed = self._committed_offset[tp]
try:
if committed is None or offset > committed:
acked_index = self._acked_index[tp]
if offset not in acked_index:
self._unacked_messages.discard(message)
acked_index.add(offset)
acked_for_tp = self._acked[tp]
acked_for_tp.append(offset)
self._n_acked += 1
return True
finally:
notify(self._waiting_for_ack)
return False
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:
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(loop=self.loop)
buffer_consumed = asyncio.Event(loop=self.loop)
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:
# 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 not 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)
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 times to
# budfer while we read.
buffer_consuming = self.loop.create_future()
try:
yield list(buffer)
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()
finally:
# Restore last behaviour of "enable_acks"
self.enable_acks = stream_enable_acks
self._processors.remove(add_to_buffer)
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
number of messages are received within the specified number of
seconds then we flush the buffer immediately.
:param 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
# We want to save events instead of values to allow for manual ack
event = self.current_event
buffer_add(cast(T_co, 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(buffer)
finally:
buffer.clear()
# code change: We want to manually ack
# 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 on_start(self) -> None:
await self.web.start_server(self.loop)
notify(self._port_open) # <- .start() can return now
