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