import collections import copy import logging import threading import time import kafka.errors as Errors from kafka.producer.producer_batch import ProducerBatch from kafka.record.memory_records import MemoryRecords, MemoryRecordsBuilder from kafka.structs import TopicPartition log = logging.getLogger(__name__) class AtomicInteger: def __init__(self, val=0): self._lock = threading.Lock() self._val = val def increment(self): with self._lock: self._val += 1 return self._val def decrement(self): with self._lock: self._val -= 1 return self._val def get(self): return self._val class RecordAccumulator: """ This class maintains a dequeue per TopicPartition that accumulates messages into MessageSets to be sent to the server. The accumulator attempts to bound memory use, and append calls will block when that memory is exhausted. Keyword Arguments: batch_size (int): Requests sent to brokers will contain multiple batches, one for each partition with data available to be sent. A small batch size will make batching less common and may reduce throughput (a batch size of zero will disable batching entirely). Default: 16384 compression_attrs (int): The compression type for all data generated by the producer. Valid values are gzip(1), snappy(2), lz4(3), or none(0). Compression is of full batches of data, so the efficacy of batching will also impact the compression ratio (more batching means better compression). Default: None. linger_ms (int): An artificial delay time to add before declaring a record batch (that isn't full) ready for sending. This allows time for more records to arrive. Setting a non-zero linger_ms will trade off some latency for potentially better throughput due to more batching (and hence fewer, larger requests). Default: 0 retry_backoff_ms (int): An artificial delay time to retry the produce request upon receiving an error. This avoids exhausting all retries in a short period of time. Default: 100 """ DEFAULT_CONFIG = { 'batch_size': 16384, 'compression_attrs': 0, 'linger_ms': 0, 'request_timeout_ms': 30000, 'delivery_timeout_ms': 120000, 'retry_backoff_ms': 100, 'transaction_manager': None, 'message_version': 2, } def __init__(self, **configs): self.config = copy.copy(self.DEFAULT_CONFIG) for key in self.config: if key in configs: self.config[key] = configs.pop(key) self._closed = False self._transaction_manager = self.config['transaction_manager'] self._flushes_in_progress = AtomicInteger() self._appends_in_progress = AtomicInteger() self._batches = collections.defaultdict(collections.deque) # TopicPartition: [ProducerBatch] self._tp_locks = {None: threading.Lock()} # TopicPartition: Lock, plus a lock to add entries self._incomplete = IncompleteProducerBatches() # The following variables should only be accessed by the sender thread, # so we don't need to protect them w/ locking. self.muted = set() self._drain_index = 0 self._next_batch_expiry_time_ms = float('inf') if self.config['delivery_timeout_ms'] < self.config['linger_ms'] + self.config['request_timeout_ms']: raise Errors.KafkaConfigurationError("Must set delivery_timeout_ms higher than linger_ms + request_timeout_ms") @property def delivery_timeout_ms(self): return self.config['delivery_timeout_ms'] @property def next_expiry_time_ms(self): return self._next_batch_expiry_time_ms def _tp_lock(self, tp): if tp not in self._tp_locks: with self._tp_locks[None]: if tp not in self._tp_locks: self._tp_locks[tp] = threading.Lock() return self._tp_locks[tp] def append(self, tp, timestamp_ms, key, value, headers, now=None, abort_on_new_batch=False): """Add a record to the accumulator, return the append result. The append result will contain the future metadata, and flag for whether the appended batch is full or a new batch is created Arguments: tp (TopicPartition): The topic/partition to which this record is being sent timestamp_ms (int): The timestamp of the record (epoch ms) key (bytes): The key for the record value (bytes): The value for the record headers (List[Tuple[str, bytes]]): The header fields for the record abort_on_new_batch (bool): KIP-480. When True, return early with ``abort_for_new_batch=True`` instead of allocating a new batch when no in-progress batch has room. Caller is expected to consult the partitioner's ``on_new_batch`` hook, re-pick the partition, and retry with ``abort_on_new_batch=False``. Returns: tuple: (future, batch_is_full, new_batch_created, abort_for_new_batch) """ if not isinstance(tp, TopicPartition): raise TypeError('not TopicPartition') if self._closed: raise Errors.IllegalStateError('RecordAccumulator is closed') now = time.monotonic() if now is None else now # We keep track of the number of appending thread to make sure we do # not miss batches in abortIncompleteBatches(). self._appends_in_progress.increment() try: with self._tp_lock(tp): # check if we have an in-progress batch dq = self._batches[tp] if dq: last = dq[-1] future = last.try_append(timestamp_ms, key, value, headers, now=now) if future is not None: batch_is_full = len(dq) > 1 or last.records.is_full() return future, batch_is_full, False, False if abort_on_new_batch: # KIP-480: don't allocate a new batch yet. Caller will # rotate the sticky partition and retry. return None, False, False, True with self._tp_lock(tp): # Need to check if producer is closed again after grabbing the # dequeue lock. if self._closed: raise Errors.IllegalStateError('RecordAccumulator is closed') if dq: last = dq[-1] future = last.try_append(timestamp_ms, key, value, headers, now=now) if future is not None: # Somebody else found us a batch, return the one we # waited for! Hopefully this doesn't happen often... batch_is_full = len(dq) > 1 or last.records.is_full() return future, batch_is_full, False, False if self._transaction_manager and self.config['message_version'] < 2: raise Errors.UnsupportedVersionError("Attempting to use idempotence with a broker which" " does not support the required message format (v2)." " The broker must be version 0.11 or later.") records = MemoryRecordsBuilder( self.config['message_version'], self.config['compression_attrs'], self.config['batch_size'] ) batch = ProducerBatch(tp, records, now=now) future = batch.try_append(timestamp_ms, key, value, headers, now=now) if not future: raise Exception() dq.append(batch) self._incomplete.add(batch) batch_is_full = len(dq) > 1 or batch.records.is_full() return future, batch_is_full, True, False finally: self._appends_in_progress.decrement() def reset_next_batch_expiry_time(self): self._next_batch_expiry_time_ms = float('inf') def maybe_update_next_batch_expiry_time(self, batch): self._next_batch_expiry_time_ms = min(self._next_batch_expiry_time_ms, batch.created * 1000 + self.delivery_timeout_ms) def expired_batches(self, now=None): """Get a list of batches which have been sitting in the accumulator too long and need to be expired.""" expired_batches = [] for tp in list(self._batches.keys()): with self._tp_lock(tp): # iterate over the batches and expire them if they have stayed # in accumulator for more than request_timeout_ms dq = self._batches[tp] while dq: batch = dq[0] if batch.has_reached_delivery_timeout(self.delivery_timeout_ms, now=now): dq.popleft() batch.records.close() expired_batches.append(batch) else: # Stop at the first batch that has not expired. self.maybe_update_next_batch_expiry_time(batch) break return expired_batches def split_and_reenqueue(self, batch, now=None): """Split an oversized batch into smaller batches and reenqueue them. When a produce request fails with MESSAGE_TOO_LARGE, this method splits the batch into two sub-batches (by record count) and enqueues them at the front of the partition's deque. The original FutureRecordMetadata objects are rebound to the new batches' futures. If the new batches are still too large, they will be split again on the next MESSAGE_TOO_LARGE response. Only supported for message_version >= 2 (DefaultRecordBatch). Arguments: batch (ProducerBatch): The oversized batch to split. Returns: int: The number of new batches created. """ now = time.monotonic() if now is None else now tp = batch.topic_partition # Roll back the partition's sequence counter to the failed batch's # base sequence. The failed batch was never committed by the broker, # so its sequence range is free to be reused by the split batches. # They will get fresh sequences assigned during drain. if self._transaction_manager: base_sequence = batch.records.base_sequence if base_sequence is not None and base_sequence != -1: self._transaction_manager.set_sequence_number(tp, base_sequence) # Read all records from the closed batch records_list = [] for record_batch in MemoryRecords(batch.records.buffer()): for record in record_batch: records_list.append(record) # Split records into two halves by count mid = (len(records_list) + 1) // 2 groups = [records_list[:mid], records_list[mid:]] new_batches = [] future_index = 0 for group in groups: if not group: continue builder = MemoryRecordsBuilder( self.config['message_version'], self.config['compression_attrs'], self.config['batch_size'], ) current_batch = ProducerBatch(tp, builder, now=now) current_batch.created = batch.created for record in group: metadata = builder.append(record.timestamp, record.key, record.value, record.headers) if metadata is None: # Record doesn't fit (extremely unlikely for split batches). # Finalize this batch and start a new one. new_batches.append(current_batch) builder = MemoryRecordsBuilder( self.config['message_version'], self.config['compression_attrs'], self.config['batch_size'], ) current_batch = ProducerBatch(tp, builder, now=now) current_batch.created = batch.created metadata = builder.append(record.timestamp, record.key, record.value, record.headers) # Rebind original future to new batch if future_index < len(batch._record_futures): original_future = batch._record_futures[future_index] original_future.rebind(current_batch.produce_future, metadata.offset) current_batch._record_futures.append(original_future) future_index += 1 new_batches.append(current_batch) # Enqueue in reverse order so first batch is at front of deque with self._tp_lock(tp): dq = self._batches[tp] for new_batch in reversed(new_batches): new_batch.attempts = batch.attempts new_batch.last_attempt = now dq.appendleft(new_batch) self._incomplete.add(new_batch) log.info("Split oversized batch for %s into %d new batches (%d total records)", tp, len(new_batches), future_index) return len(new_batches) def reenqueue(self, batch, now=None): """ Re-enqueue the given record batch in the accumulator. In Sender._complete_batch method, we check whether the batch has reached delivery_timeout_ms or not. Hence we do not do the delivery timeout check here. """ batch.retry(now=now) with self._tp_lock(batch.topic_partition): dq = self._batches[batch.topic_partition] dq.appendleft(batch) def ready(self, cluster, now=None): """ Get a list of nodes whose partitions are ready to be sent, and the earliest time at which any non-sendable partition will be ready; Also return the flag for whether there are any unknown leaders for the accumulated partition batches. A destination node is ready to send if: * There is at least one partition that is not backing off its send * and those partitions are not muted (to prevent reordering if max_in_flight_requests_per_connection is set to 1) * and any of the following are true: * The record set is full * The record set has sat in the accumulator for at least linger_ms milliseconds * The accumulator is out of memory and threads are blocking waiting for data (in this case all partitions are immediately considered ready). * The accumulator has been closed Arguments: cluster (ClusterMetadata): Returns: tuple: ready_nodes (set): node_ids that have ready batches next_ready_check (float): secs until next ready after backoff unknown_leaders_exist (bool): True if metadata refresh needed """ ready_nodes = set() next_ready_check = 9999999.99 unknown_leaders_exist = False now = time.monotonic() if now is None else now # several threads are accessing self._batches -- to simplify # concurrent access, we iterate over a snapshot of partitions # and lock each partition separately as needed partitions = list(self._batches.keys()) for tp in partitions: leader = cluster.leader_for_partition(tp) if leader is None or leader == -1: unknown_leaders_exist = True continue elif leader in ready_nodes: continue elif tp in self.muted: continue with self._tp_lock(tp): dq = self._batches[tp] if not dq: continue batch = dq[0] retry_backoff = self.config['retry_backoff_ms'] / 1000 linger = self.config['linger_ms'] / 1000 backing_off = bool(batch.attempts > 0 and (batch.last_attempt + retry_backoff) > now) waited_time = now - batch.last_attempt time_to_wait = retry_backoff if backing_off else linger time_left = max(time_to_wait - waited_time, 0) full = bool(len(dq) > 1 or batch.records.is_full()) expired = bool(waited_time >= time_to_wait) sendable = (full or expired or self._closed or self.flush_in_progress()) if sendable and not backing_off: ready_nodes.add(leader) else: # Note that this results in a conservative estimate since # an un-sendable partition may have a leader that will # later be found to have sendable data. However, this is # good enough since we'll just wake up and then sleep again # for the remaining time. next_ready_check = min(time_left, next_ready_check) return ready_nodes, next_ready_check, unknown_leaders_exist def has_undrained(self): """Check whether there are any batches which haven't been drained""" for tp in list(self._batches.keys()): with self._tp_lock(tp): dq = self._batches[tp] if len(dq): return True return False def _should_stop_drain_batches_for_partition(self, first, tp): if self._transaction_manager: if not self._transaction_manager.is_send_to_partition_allowed(tp): return True if not self._transaction_manager.producer_id_and_epoch.is_valid: # we cannot send the batch until we have refreshed the PID log.debug("Waiting to send ready batches because transaction producer id is not valid") return True return False def drain_batches_for_one_node(self, cluster, node_id, max_size, now=None): now = time.monotonic() if now is None else now size = 0 ready = [] partitions = list(cluster.partitions_for_broker(node_id)) if not partitions: return ready # to make starvation less likely this loop doesn't start at 0 self._drain_index %= len(partitions) start = None while start != self._drain_index: tp = partitions[self._drain_index] if start is None: start = self._drain_index self._drain_index += 1 self._drain_index %= len(partitions) # Only proceed if the partition has no in-flight batches. if tp in self.muted: continue if tp not in self._batches: continue with self._tp_lock(tp): dq = self._batches[tp] if len(dq) == 0: continue first = dq[0] backoff = bool(first.attempts > 0 and first.last_attempt + self.config['retry_backoff_ms'] / 1000 > now) # Only drain the batch if it is not during backoff if backoff: continue if (size + first.records.size_in_bytes() > max_size and len(ready) > 0): # there is a rare case that a single batch # size is larger than the request size due # to compression; in this case we will # still eventually send this batch in a # single request break else: if self._should_stop_drain_batches_for_partition(first, tp): break batch = dq.popleft() if self._transaction_manager and not batch.in_retry(): # If the batch is in retry, then we should not change the pid and # sequence number, since this may introduce duplicates. In particular, # the previous attempt may actually have been accepted, and if we change # the pid and sequence here, this attempt will also be accepted, causing # a duplicate. sequence_number = self._transaction_manager.sequence_number(batch.topic_partition) log.debug("Dest: %s: %s producer_id=%s epoch=%s sequence=%s", node_id, batch.topic_partition, self._transaction_manager.producer_id_and_epoch.producer_id, self._transaction_manager.producer_id_and_epoch.epoch, sequence_number) batch.records.set_producer_state( self._transaction_manager.producer_id_and_epoch.producer_id, self._transaction_manager.producer_id_and_epoch.epoch, sequence_number, self._transaction_manager.is_transactional() ) # Increment sequence now so subsequent in-flight batches # for the same partition get the correct next sequence. self._transaction_manager.increment_sequence_number( batch.topic_partition, batch.records.next_offset()) batch.records.close() size += batch.records.size_in_bytes() ready.append(batch) batch.drained = now return ready def drain(self, cluster, nodes, max_size, now=None): """ Drain all the data for the given nodes and collate them into a list of batches that will fit within the specified size on a per-node basis. This method attempts to avoid choosing the same topic-node repeatedly. Arguments: cluster (ClusterMetadata): The current cluster metadata nodes (list): list of node_ids to drain max_size (int): maximum number of bytes to drain Returns: dict: {node_id: list of ProducerBatch} with total size less than the requested max_size. """ if not nodes: return {} now = time.monotonic() if now is None else now batches = {} for node_id in nodes: batches[node_id] = self.drain_batches_for_one_node(cluster, node_id, max_size, now=now) return batches def deallocate(self, batch): """Deallocate the record batch.""" self._incomplete.remove(batch) def flush_in_progress(self): """Are there any threads currently waiting on a flush?""" return self._flushes_in_progress.get() > 0 def begin_flush(self): """ Initiate the flushing of data from the accumulator...this makes all requests immediately ready """ self._flushes_in_progress.increment() def await_flush_completion(self, timeout=None): """ Mark all partitions as ready to send and block until the send is complete """ try: for batch in self._incomplete.all(): log.debug('Waiting on produce to %s', batch.produce_future.topic_partition) if not batch.produce_future.wait(timeout=timeout): raise Errors.KafkaTimeoutError('Timeout waiting for future') if not batch.produce_future.is_done: raise Errors.UnknownError('Future not done') if batch.produce_future.failed(): log.warning(batch.produce_future.exception) finally: self._flushes_in_progress.decrement() @property def has_incomplete(self): return bool(self._incomplete) def abort_incomplete_batches(self): """ This function is only called when sender is closed forcefully. It will fail all the incomplete batches and return. """ # We need to keep aborting the incomplete batch until no thread is trying to append to # 1. Avoid losing batches. # 2. Free up memory in case appending threads are blocked on buffer full. # This is a tight loop but should be able to get through very quickly. error = Errors.IllegalStateError("Producer is closed forcefully.") while True: self.abort_batches(error) if not self._appends_in_progress.get(): break # After this point, no thread will append any messages because they will see the close # flag set. We need to do the last abort after no thread was appending in case the there was a new # batch appended by the last appending thread. self.abort_batches(error) self._batches.clear() def abort_batches(self, error): """Abort every incomplete batch, including in-flight (drained but not-yet-acked) ones. Use for fatal-error / force-close paths where the user's pending futures must resolve immediately rather than hang waiting for broker responses that aren't coming.""" for batch in self._incomplete.all(): tp = batch.topic_partition with self._tp_lock(tp): batch.records.close() # Drained batches were popleft()'d out of _batches[tp] -- only # remove if still present (matches Java's LinkedList.remove, # which returns false rather than raising). try: self._batches[tp].remove(batch) except ValueError: pass batch.abort(error) self.deallocate(batch) def abort_undrained_batches(self, error): for batch in self._incomplete.all(): tp = batch.topic_partition with self._tp_lock(tp): aborted = False # Skip in-flight batches (already drained, awaiting broker # response): drain() popped them from _batches[tp], so # .remove() would ValueError; and we want their futures to # resolve via the broker response, not be cancelled locally. if not batch.is_done and batch.drained is None: aborted = True batch.records.close() self._batches[tp].remove(batch) if aborted: batch.abort(error) self.deallocate(batch) def close(self): """Close this accumulator and force all the record buffers to be drained.""" self._closed = True class IncompleteProducerBatches: """A threadsafe helper class to hold ProducerBatches that haven't been ack'd yet""" def __init__(self): self._incomplete = set() self._lock = threading.Lock() def add(self, batch): with self._lock: self._incomplete.add(batch) def remove(self, batch): with self._lock: try: self._incomplete.remove(batch) except KeyError: pass def all(self): with self._lock: return list(self._incomplete) def __bool__(self): return bool(self._incomplete) __nonzero__ = __bool__