async def add_batch(self, builder, tp, timeout):
"""Add BatchBuilder to queue by topic-partition.
Arguments:
builder (BatchBuilder): batch object to enqueue.
tp (TopicPartition): topic and partition to enqueue this batch for.
timeout (int): time in seconds to wait for a free slot in the batch
queue.
Returns:
MessageBatch: delivery wrapper around the BatchBuilder object.
Raises:
aiokafka.errors.ProducerClosed: the accumulator has already been
closed and flushed.
aiokafka.errors.KafkaTimeoutError: the batch could not be added
within the specified timeout.
"""
if self._closed:
raise ProducerClosed()
if self._exception is not None:
raise copy.copy(self._exception)
start = self._loop.time()
while timeout > 0:
pending = self._batches.get(tp)
if pending:
await pending[-1].wait_drain(timeout=timeout)
timeout -= self._loop.time() - start
else:
batch = self._append_batch(builder, tp)
return asyncio.shield(batch.future, loop=self._loop)
raise KafkaTimeoutError()
async def add_message(
self, tp, key, value, timeout, timestamp_ms=None,
headers=[]
):
""" Add message to batch by topic-partition
If batch is already full this method waits (`timeout` seconds maximum)
until batch is drained by send task
"""
if self._closed:
# this can happen when producer is closing but try to send some
# messages in async task
raise ProducerClosed()
if self._exception is not None:
raise copy.copy(self._exception)
pending_batches = self._batches.get(tp)
if not pending_batches:
builder = self.create_builder()
batch = self._append_batch(builder, tp)
else:
batch = pending_batches[-1]
future = batch.append(key, value, timestamp_ms, headers=headers)
if future is None:
# Batch is full, can't append data atm,
# waiting until batch per topic-partition is drained
start = self._loop.time()
await batch.wait_drain(timeout)
timeout -= self._loop.time() - start
if timeout <= 0:
raise KafkaTimeoutError()
return (await self.add_message(
tp, key, value, timeout, timestamp_ms))
return future
def add_message(self, tp, key, value, timeout):
""" Add message to batch by topic-partition
If batch is already full this method waits (`ttl` seconds maximum)
until batch is drained by send task
"""
if self._closed:
# this can happen when producer is closing but try to send some
# messages in async task
raise ProducerClosed()
batch = self._batches.get(tp)
if not batch:
batch = MessageBatch(tp, self._batch_size, self._compression_type,
self._batch_ttl, self._api_version,
self._loop)
self._batches[tp] = batch
if not self._wait_data_future.done():
# Wakeup sender task if it waits for data
self._wait_data_future.set_result(None)
future = batch.append(key, value)
if future is None:
# Batch is full, can't append data atm,
# waiting until batch per topic-partition is drained
start = self._loop.time()
yield from asyncio.wait([batch.wait_drain()],
timeout=timeout,
loop=self._loop)
timeout -= self._loop.time() - start
if timeout <= 0:
raise KafkaTimeoutError()
return (yield from self.add_message(tp, key, value, timeout))
return future
async def _retrieve_offsets(self, timestamps, timeout_ms=float("inf")):
""" Fetch offset for each partition passed in ``timestamps`` map.
Blocks until offsets are obtained, a non-retriable exception is raised
or ``timeout_ms`` passed.
Arguments:
timestamps: {TopicPartition: int} dict with timestamps to fetch
offsets by. -1 for the latest available, -2 for the earliest
available. Otherwise timestamp is treated as epoch
milliseconds.
Returns:
{TopicPartition: (int, int)}: Mapping of partition to
retrieved offset and timestamp. If offset does not exist for
the provided timestamp, that partition will be missing from
this mapping.
"""
if not timestamps:
return {}
timeout = timeout_ms / 1000
start_time = self._loop.time()
remaining = timeout
while True:
try:
offsets = await asyncio.wait_for(
self._proc_offset_requests(timestamps),
timeout=None if remaining == float("inf") else remaining,
loop=self._loop
)
except asyncio.TimeoutError:
break
except Errors.KafkaError as error:
if not error.retriable:
raise error
if error.invalid_metadata:
self._client.force_metadata_update()
elapsed = self._loop.time() - start_time
remaining = max(0, remaining - elapsed)
if remaining < self._retry_backoff:
break
await asyncio.sleep(self._retry_backoff, loop=self._loop)
else:
return offsets
raise KafkaTimeoutError(
"Failed to get offsets by times in %s ms" % timeout_ms)
async def _retrieve_offsets(self, timestamps, timeout_ms=None):
""" Fetch offset for each partition passed in ``timestamps`` map.
Blocks until offsets are obtained, a non-retriable exception is raised
or ``timeout_ms`` passed.
Arguments:
timestamps: {TopicPartition: int} dict with timestamps to fetch
offsets by. -1 for the latest available, -2 for the earliest
available. Otherwise timestamp is treated as epoch
milliseconds.
Returns:
{TopicPartition: (int, int)}: Mapping of partition to
retrieved offset and timestamp. If offset does not exist for
the provided timestamp, that partition will be missing from
this mapping.
"""
if not timestamps:
return {}
timeout = None if timeout_ms is None else timeout_ms / 1000
try:
async with async_timeout.timeout(timeout):
while True:
try:
offsets = await self._proc_offset_requests(timestamps)
except Errors.KafkaError as error:
if not error.retriable:
raise error
if error.invalid_metadata:
self._client.force_metadata_update()
await asyncio.sleep(self._retry_backoff)
else:
return offsets
except asyncio.TimeoutError:
raise KafkaTimeoutError("Failed to get offsets by times in %s ms" %
timeout_ms)
def send(self, node_id, request):
"""Send a request to a specific node.
Arguments:
node_id (int): destination node
request (Struct): request object (not-encoded)
Raises:
kafka.common.KafkaTimeoutError
kafka.common.NodeNotReadyError
kafka.common.ConnectionError
kafka.common.CorrelationIdError
Returns:
Future: resolves to Response struct
"""
if not (yield from self.ready(node_id)):
raise NodeNotReadyError(
"Attempt to send a request to node"
" which is not ready (node id {}).".format(node_id))
# Every request gets a response, except one special case:
expect_response = True
if isinstance(request, tuple(ProduceRequest)) and \
request.required_acks == 0:
expect_response = False
future = self._conns[node_id].send(request,
expect_response=expect_response)
try:
result = yield from future
except asyncio.TimeoutError:
# close connection so it is renewed in next request
self._conns[node_id].close()
raise KafkaTimeoutError()
else:
return result
async def mocked_send(nodeid, req):
raise KafkaTimeoutError()
