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
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()
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
