init

venv/lib/python3.11/site-packages/kafka/producer/__init__.py

@@ -0,0 +1,7 @@
+from __future__ import absolute_import
+
+from kafka.producer.kafka import KafkaProducer
+
+__all__ = [
+    'KafkaProducer'
+]

venv/lib/python3.11/site-packages/kafka/producer/future.py

@@ -0,0 +1,71 @@
+from __future__ import absolute_import
+
+import collections
+import threading
+
+from kafka import errors as Errors
+from kafka.future import Future
+
+
+class FutureProduceResult(Future):
+    def __init__(self, topic_partition):
+        super(FutureProduceResult, self).__init__()
+        self.topic_partition = topic_partition
+        self._latch = threading.Event()
+
+    def success(self, value):
+        ret = super(FutureProduceResult, self).success(value)
+        self._latch.set()
+        return ret
+
+    def failure(self, error):
+        ret = super(FutureProduceResult, self).failure(error)
+        self._latch.set()
+        return ret
+
+    def wait(self, timeout=None):
+        # wait() on python2.6 returns None instead of the flag value
+        return self._latch.wait(timeout) or self._latch.is_set()
+
+
+class FutureRecordMetadata(Future):
+    def __init__(self, produce_future, relative_offset, timestamp_ms, checksum, serialized_key_size, serialized_value_size, serialized_header_size):
+        super(FutureRecordMetadata, self).__init__()
+        self._produce_future = produce_future
+        # packing args as a tuple is a minor speed optimization
+        self.args = (relative_offset, timestamp_ms, checksum, serialized_key_size, serialized_value_size, serialized_header_size)
+        produce_future.add_callback(self._produce_success)
+        produce_future.add_errback(self.failure)
+
+    def _produce_success(self, offset_and_timestamp):
+        offset, produce_timestamp_ms = offset_and_timestamp
+
+        # Unpacking from args tuple is minor speed optimization
+        (relative_offset, timestamp_ms, checksum,
+         serialized_key_size, serialized_value_size, serialized_header_size) = self.args
+
+        # None is when Broker does not support the API (<0.10) and
+        # -1 is when the broker is configured for CREATE_TIME timestamps
+        if produce_timestamp_ms is not None and produce_timestamp_ms != -1:
+            timestamp_ms = produce_timestamp_ms
+        if offset != -1 and relative_offset is not None:
+            offset += relative_offset
+        tp = self._produce_future.topic_partition
+        metadata = RecordMetadata(tp[0], tp[1], tp, offset, timestamp_ms,
+                                  checksum, serialized_key_size,
+                                  serialized_value_size, serialized_header_size)
+        self.success(metadata)
+
+    def get(self, timeout=None):
+        if not self.is_done and not self._produce_future.wait(timeout):
+            raise Errors.KafkaTimeoutError(
+                "Timeout after waiting for %s secs." % (timeout,))
+        assert self.is_done
+        if self.failed():
+            raise self.exception # pylint: disable-msg=raising-bad-type
+        return self.value
+
+
+RecordMetadata = collections.namedtuple(
+    'RecordMetadata', ['topic', 'partition', 'topic_partition', 'offset', 'timestamp',
+                       'checksum', 'serialized_key_size', 'serialized_value_size', 'serialized_header_size'])

venv/lib/python3.11/site-packages/kafka/producer/record_accumulator.py

@@ -0,0 +1,674 @@
+from __future__ import absolute_import, division
+
+import collections
+import copy
+import logging
+import threading
+import time
+
+try:
+    # enum in stdlib as of py3.4
+    from enum import IntEnum  # pylint: disable=import-error
+except ImportError:
+    # vendored backport module
+    from kafka.vendor.enum34 import IntEnum
+
+import kafka.errors as Errors
+from kafka.producer.future import FutureRecordMetadata, FutureProduceResult
+from kafka.record.memory_records import MemoryRecordsBuilder
+from kafka.structs import TopicPartition
+
+
+log = logging.getLogger(__name__)
+
+
+class AtomicInteger(object):
+    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 FinalState(IntEnum):
+    ABORTED = 0
+    FAILED = 1
+    SUCCEEDED = 2
+
+
+class ProducerBatch(object):
+    def __init__(self, tp, records, now=None):
+        now = time.time() if now is None else now
+        self.max_record_size = 0
+        self.created = now
+        self.drained = None
+        self.attempts = 0
+        self.last_attempt = now
+        self.last_append = now
+        self.records = records
+        self.topic_partition = tp
+        self.produce_future = FutureProduceResult(tp)
+        self._retry = False
+        self._final_state = None
+
+    @property
+    def final_state(self):
+        return self._final_state
+
+    @property
+    def record_count(self):
+        return self.records.next_offset()
+
+    @property
+    def producer_id(self):
+        return self.records.producer_id if self.records else None
+
+    @property
+    def producer_epoch(self):
+        return self.records.producer_epoch if self.records else None
+
+    @property
+    def has_sequence(self):
+        return self.records.has_sequence if self.records else False
+
+    def try_append(self, timestamp_ms, key, value, headers, now=None):
+        metadata = self.records.append(timestamp_ms, key, value, headers)
+        if metadata is None:
+            return None
+
+        now = time.time() if now is None else now
+        self.max_record_size = max(self.max_record_size, metadata.size)
+        self.last_append = now
+        future = FutureRecordMetadata(
+            self.produce_future,
+            metadata.offset,
+            metadata.timestamp,
+            metadata.crc,
+            len(key) if key is not None else -1,
+            len(value) if value is not None else -1,
+            sum(len(h_key.encode("utf-8")) + len(h_val) for h_key, h_val in headers) if headers else -1)
+        return future
+
+    def abort(self, exception):
+        """Abort the batch and complete the future and callbacks."""
+        if self._final_state is not None:
+            raise Errors.IllegalStateError("Batch has already been completed in final state: %s" % self._final_state)
+        self._final_state = FinalState.ABORTED
+
+        log.debug("Aborting batch for partition %s: %s", self.topic_partition, exception)
+        self._complete_future(-1, -1, exception)
+
+    def done(self, base_offset=None, timestamp_ms=None, exception=None):
+        """
+        Finalize the state of a batch. Final state, once set, is immutable. This function may be called
+        once or twice on a batch. It may be called twice if
+            1. An inflight batch expires before a response from the broker is received. The batch's final
+            state is set to FAILED. But it could succeed on the broker and second time around batch.done() may
+            try to set SUCCEEDED final state.
+
+            2. If a transaction abortion happens or if the producer is closed forcefully, the final state is
+            ABORTED but again it could succeed if broker responds with a success.
+
+        Attempted transitions from [FAILED | ABORTED] --> SUCCEEDED are logged.
+        Attempted transitions from one failure state to the same or a different failed state are ignored.
+        Attempted transitions from SUCCEEDED to the same or a failed state throw an exception.
+        """
+        final_state = FinalState.SUCCEEDED if exception is None else FinalState.FAILED
+        if self._final_state is None:
+            self._final_state = final_state
+            if final_state is FinalState.SUCCEEDED:
+                log.debug("Successfully produced messages to %s with base offset %s", self.topic_partition, base_offset)
+            else:
+                log.warning("Failed to produce messages to topic-partition %s with base offset %s: %s",
+                            self.topic_partition, base_offset, exception)
+            self._complete_future(base_offset, timestamp_ms, exception)
+            return True
+
+        elif self._final_state is not FinalState.SUCCEEDED:
+            if final_state is FinalState.SUCCEEDED:
+                # Log if a previously unsuccessful batch succeeded later on.
+                log.debug("ProduceResponse returned %s for %s after batch with base offset %s had already been %s.",
+                          final_state, self.topic_partition, base_offset, self._final_state)
+            else:
+                # FAILED --> FAILED and ABORTED --> FAILED transitions are ignored.
+                log.debug("Ignored state transition %s -> %s for %s batch with base offset %s",
+                          self._final_state, final_state, self.topic_partition, base_offset)
+        else:
+            # A SUCCESSFUL batch must not attempt another state change.
+            raise Errors.IllegalStateError("A %s batch must not attempt another state change to %s" % (self._final_state, final_state))
+        return False
+
+    def _complete_future(self, base_offset, timestamp_ms, exception):
+        if self.produce_future.is_done:
+            raise Errors.IllegalStateError('Batch is already closed!')
+        elif exception is None:
+            self.produce_future.success((base_offset, timestamp_ms))
+        else:
+            self.produce_future.failure(exception)
+
+    def has_reached_delivery_timeout(self, delivery_timeout_ms, now=None):
+        now = time.time() if now is None else now
+        return delivery_timeout_ms / 1000 <= now - self.created
+
+    def in_retry(self):
+        return self._retry
+
+    def retry(self, now=None):
+        now = time.time() if now is None else now
+        self._retry = True
+        self.attempts += 1
+        self.last_attempt = now
+        self.last_append = now
+
+    @property
+    def is_done(self):
+        return self.produce_future.is_done
+
+    def __str__(self):
+        return 'ProducerBatch(topic_partition=%s, record_count=%d)' % (
+            self.topic_partition, self.records.next_offset())
+
+
+class RecordAccumulator(object):
+    """
+    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):
+        """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
+
+        Returns:
+            tuple: (future, batch_is_full, new_batch_created)
+        """
+        assert isinstance(tp, TopicPartition), 'not TopicPartition'
+        assert not self._closed, 'RecordAccumulator is closed'
+        now = time.time() 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
+
+            with self._tp_lock(tp):
+                # Need to check if producer is closed again after grabbing the
+                # dequeue lock.
+                assert not self._closed, '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
+
+                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
+        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 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.time() 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.time() 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()
+                        )
+                    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.time() 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):
+        """Go through incomplete batches and abort them."""
+        for batch in self._incomplete.all():
+            tp = batch.topic_partition
+            # Close the batch before aborting
+            with self._tp_lock(tp):
+                batch.records.close()
+                self._batches[tp].remove(batch)
+            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
+                if not batch.is_done:
+                    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(object):
+    """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__

venv/lib/python3.11/site-packages/kafka/producer/sender.py

@@ -0,0 +1,763 @@
+from __future__ import absolute_import, division
+
+import collections
+import copy
+import heapq
+import logging
+import threading
+import time
+
+from kafka.vendor import six
+
+from kafka import errors as Errors
+from kafka.metrics.measurable import AnonMeasurable
+from kafka.metrics.stats import Avg, Max, Rate
+from kafka.producer.transaction_manager import ProducerIdAndEpoch
+from kafka.protocol.init_producer_id import InitProducerIdRequest
+from kafka.protocol.produce import ProduceRequest
+from kafka.structs import TopicPartition
+from kafka.version import __version__
+
+log = logging.getLogger(__name__)
+
+
+class Sender(threading.Thread):
+    """
+    The background thread that handles the sending of produce requests to the
+    Kafka cluster. This thread makes metadata requests to renew its view of the
+    cluster and then sends produce requests to the appropriate nodes.
+    """
+    DEFAULT_CONFIG = {
+        'max_request_size': 1048576,
+        'acks': 1,
+        'retries': float('inf'),
+        'request_timeout_ms': 30000,
+        'retry_backoff_ms': 100,
+        'metrics': None,
+        'guarantee_message_order': False,
+        'transaction_manager': None,
+        'transactional_id': None,
+        'transaction_timeout_ms': 60000,
+        'client_id': 'kafka-python-' + __version__,
+    }
+
+    def __init__(self, client, metadata, accumulator, **configs):
+        super(Sender, self).__init__()
+        self.config = copy.copy(self.DEFAULT_CONFIG)
+        for key in self.config:
+            if key in configs:
+                self.config[key] = configs.pop(key)
+
+        self.name = self.config['client_id'] + '-network-thread'
+        self._client = client
+        self._accumulator = accumulator
+        self._metadata = client.cluster
+        self._running = True
+        self._force_close = False
+        self._topics_to_add = set()
+        if self.config['metrics']:
+            self._sensors = SenderMetrics(self.config['metrics'], self._client, self._metadata)
+        else:
+            self._sensors = None
+        self._transaction_manager = self.config['transaction_manager']
+        # A per-partition queue of batches ordered by creation time for tracking the in-flight batches
+        self._in_flight_batches = collections.defaultdict(list)
+
+    def _maybe_remove_from_inflight_batches(self, batch):
+        try:
+            queue = self._in_flight_batches[batch.topic_partition]
+        except KeyError:
+            return
+        try:
+            idx = queue.index((batch.created, batch))
+        except ValueError:
+            return
+        # https://stackoverflow.com/questions/10162679/python-delete-element-from-heap
+        queue[idx] = queue[-1]
+        queue.pop()
+        heapq.heapify(queue)
+
+    def _get_expired_inflight_batches(self, now=None):
+        """Get the in-flight batches that has reached delivery timeout."""
+        expired_batches = []
+        to_remove = []
+        for tp, queue in six.iteritems(self._in_flight_batches):
+            while queue:
+                _created_at, batch = queue[0]
+                if batch.has_reached_delivery_timeout(self._accumulator.delivery_timeout_ms):
+                    heapq.heappop(queue)
+                    if batch.final_state is None:
+                        expired_batches.append(batch)
+                    else:
+                        raise Errors.IllegalStateError("%s batch created at %s gets unexpected final state %s" % (batch.topic_partition, batch.created, batch.final_state))
+                else:
+                    self._accumulator.maybe_update_next_batch_expiry_time(batch)
+                    break
+            else:
+                # Avoid mutating in_flight_batches during iteration
+                to_remove.append(tp)
+        for tp in to_remove:
+            del self._in_flight_batches[tp]
+        return expired_batches
+
+    def run(self):
+        """The main run loop for the sender thread."""
+        log.debug("%s: Starting Kafka producer I/O thread.", str(self))
+
+        # main loop, runs until close is called
+        while self._running:
+            try:
+                self.run_once()
+            except Exception:
+                log.exception("%s: Uncaught error in kafka producer I/O thread", str(self))
+
+        log.debug("%s: Beginning shutdown of Kafka producer I/O thread, sending"
+                  " remaining records.", str(self))
+
+        # okay we stopped accepting requests but there may still be
+        # requests in the accumulator or waiting for acknowledgment,
+        # wait until these are completed.
+        while (not self._force_close
+               and (self._accumulator.has_undrained()
+                    or self._client.in_flight_request_count() > 0)):
+            try:
+                self.run_once()
+            except Exception:
+                log.exception("%s: Uncaught error in kafka producer I/O thread", str(self))
+
+        if self._force_close:
+            # We need to fail all the incomplete batches and wake up the
+            # threads waiting on the futures.
+            self._accumulator.abort_incomplete_batches()
+
+        try:
+            self._client.close()
+        except Exception:
+            log.exception("%s: Failed to close network client", str(self))
+
+        log.debug("%s: Shutdown of Kafka producer I/O thread has completed.", str(self))
+
+    def run_once(self):
+        """Run a single iteration of sending."""
+        while self._topics_to_add:
+            self._client.add_topic(self._topics_to_add.pop())
+
+        if self._transaction_manager:
+            try:
+                if not self._transaction_manager.is_transactional():
+                    # this is an idempotent producer, so make sure we have a producer id
+                    self._maybe_wait_for_producer_id()
+                elif self._transaction_manager.has_in_flight_transactional_request() or self._maybe_send_transactional_request():
+                    # as long as there are outstanding transactional requests, we simply wait for them to return
+                    self._client.poll(timeout_ms=self.config['retry_backoff_ms'])
+                    return
+
+                # do not continue sending if the transaction manager is in a failed state or if there
+                # is no producer id (for the idempotent case).
+                if self._transaction_manager.has_fatal_error() or not self._transaction_manager.has_producer_id():
+                    last_error = self._transaction_manager.last_error
+                    if last_error is not None:
+                        self._maybe_abort_batches(last_error)
+                    self._client.poll(timeout_ms=self.config['retry_backoff_ms'])
+                    return
+                elif self._transaction_manager.has_abortable_error():
+                    self._accumulator.abort_undrained_batches(self._transaction_manager.last_error)
+
+            except Errors.SaslAuthenticationFailedError as e:
+                # This is already logged as error, but propagated here to perform any clean ups.
+                log.debug("%s: Authentication exception while processing transactional request: %s", str(self), e)
+                self._transaction_manager.authentication_failed(e)
+
+        poll_timeout_ms = self._send_producer_data()
+        self._client.poll(timeout_ms=poll_timeout_ms)
+
+    def _send_producer_data(self, now=None):
+        now = time.time() if now is None else now
+        # get the list of partitions with data ready to send
+        result = self._accumulator.ready(self._metadata, now=now)
+        ready_nodes, next_ready_check_delay, unknown_leaders_exist = result
+
+        # if there are any partitions whose leaders are not known yet, force
+        # metadata update
+        if unknown_leaders_exist:
+            log.debug('%s: Unknown leaders exist, requesting metadata update', str(self))
+            self._metadata.request_update()
+
+        # remove any nodes we aren't ready to send to
+        not_ready_timeout_ms = float('inf')
+        for node in list(ready_nodes):
+            if not self._client.is_ready(node):
+                node_delay_ms = self._client.connection_delay(node)
+                log.debug('%s: Node %s not ready; delaying produce of accumulated batch (%f ms)', str(self), node, node_delay_ms)
+                self._client.maybe_connect(node, wakeup=False)
+                ready_nodes.remove(node)
+                not_ready_timeout_ms = min(not_ready_timeout_ms, node_delay_ms)
+
+        # create produce requests
+        batches_by_node = self._accumulator.drain(
+            self._metadata, ready_nodes, self.config['max_request_size'], now=now)
+
+        for batch_list in six.itervalues(batches_by_node):
+            for batch in batch_list:
+                item = (batch.created, batch)
+                queue = self._in_flight_batches[batch.topic_partition]
+                heapq.heappush(queue, item)
+
+        if self.config['guarantee_message_order']:
+            # Mute all the partitions drained
+            for batch_list in six.itervalues(batches_by_node):
+                for batch in batch_list:
+                    self._accumulator.muted.add(batch.topic_partition)
+
+        self._accumulator.reset_next_batch_expiry_time()
+        expired_batches = self._accumulator.expired_batches(now=now)
+        expired_batches.extend(self._get_expired_inflight_batches(now=now))
+
+        if expired_batches:
+            log.debug("%s: Expired %s batches in accumulator", str(self), len(expired_batches))
+
+        # Reset the producer_id if an expired batch has previously been sent to the broker.
+        # See the documentation of `TransactionState.reset_producer_id` to understand why
+        # we need to reset the producer id here.
+        if self._transaction_manager and any([batch.in_retry() for batch in expired_batches]):
+            needs_transaction_state_reset = True
+        else:
+            needs_transaction_state_reset = False
+
+        for expired_batch in expired_batches:
+            error = Errors.KafkaTimeoutError(
+                "Expiring %d record(s) for %s: %s ms has passed since batch creation" % (
+                    expired_batch.record_count, expired_batch.topic_partition,
+                    int((time.time() - expired_batch.created) * 1000)))
+            self._fail_batch(expired_batch, error, base_offset=-1)
+
+        if self._sensors:
+            self._sensors.update_produce_request_metrics(batches_by_node)
+
+        if needs_transaction_state_reset:
+            self._transaction_manager.reset_producer_id()
+            return 0
+
+        requests = self._create_produce_requests(batches_by_node)
+        # If we have any nodes that are ready to send + have sendable data,
+        # poll with 0 timeout so this can immediately loop and try sending more
+        # data. Otherwise, the timeout will be the smaller value between next
+        # batch expiry time, and the delay time for checking data availability.
+        # Note that the nodes may have data that isn't yet sendable due to
+        # lingering, backing off, etc. This specifically does not include nodes with
+        # sendable data that aren't ready to send since they would cause busy
+        # looping.
+        poll_timeout_ms = min(next_ready_check_delay * 1000,
+                              not_ready_timeout_ms,
+                              self._accumulator.next_expiry_time_ms - now * 1000)
+        if poll_timeout_ms < 0:
+            poll_timeout_ms = 0
+
+        if ready_nodes:
+            log.debug("%s: Nodes with data ready to send: %s", str(self), ready_nodes) # trace
+            log.debug("%s: Created %d produce requests: %s", str(self), len(requests), requests) # trace
+            # if some partitions are already ready to be sent, the select time
+            # would be 0; otherwise if some partition already has some data
+            # accumulated but not ready yet, the select time will be the time
+            # difference between now and its linger expiry time; otherwise the
+            # select time will be the time difference between now and the
+            # metadata expiry time
+            poll_timeout_ms = 0
+
+        for node_id, request in six.iteritems(requests):
+            batches = batches_by_node[node_id]
+            log.debug('%s: Sending Produce Request: %r', str(self), request)
+            (self._client.send(node_id, request, wakeup=False)
+                 .add_callback(
+                     self._handle_produce_response, node_id, time.time(), batches)
+                 .add_errback(
+                     self._failed_produce, batches, node_id))
+        return poll_timeout_ms
+
+    def _maybe_send_transactional_request(self):
+        if self._transaction_manager.is_completing() and self._accumulator.has_incomplete:
+            if self._transaction_manager.is_aborting():
+                self._accumulator.abort_undrained_batches(Errors.KafkaError("Failing batch since transaction was aborted"))
+            # There may still be requests left which are being retried. Since we do not know whether they had
+            # been successfully appended to the broker log, we must resend them until their final status is clear.
+            # If they had been appended and we did not receive the error, then our sequence number would no longer
+            # be correct which would lead to an OutOfSequenceNumberError.
+            if not self._accumulator.flush_in_progress():
+                self._accumulator.begin_flush()
+
+        next_request_handler = self._transaction_manager.next_request_handler(self._accumulator.has_incomplete)
+        if next_request_handler is None:
+            return False
+
+        log.debug("%s: Sending transactional request %s", str(self), next_request_handler.request)
+        while not self._force_close:
+            target_node = None
+            try:
+                if next_request_handler.needs_coordinator():
+                    target_node = self._transaction_manager.coordinator(next_request_handler.coordinator_type)
+                    if target_node is None:
+                        self._transaction_manager.lookup_coordinator_for_request(next_request_handler)
+                        break
+                    elif not self._client.await_ready(target_node, timeout_ms=self.config['request_timeout_ms']):
+                        self._transaction_manager.lookup_coordinator_for_request(next_request_handler)
+                        target_node = None
+                        break
+                else:
+                    target_node = self._client.least_loaded_node()
+                    if target_node is not None and not self._client.await_ready(target_node, timeout_ms=self.config['request_timeout_ms']):
+                        target_node = None
+
+                if target_node is not None:
+                    if next_request_handler.is_retry:
+                        time.sleep(self.config['retry_backoff_ms'] / 1000)
+                    txn_correlation_id = self._transaction_manager.next_in_flight_request_correlation_id()
+                    future = self._client.send(target_node, next_request_handler.request)
+                    future.add_both(next_request_handler.on_complete, txn_correlation_id)
+                    return True
+
+            except Exception as e:
+                log.warn("%s: Got an exception when trying to find a node to send a transactional request to. Going to back off and retry: %s", str(self), e)
+                if next_request_handler.needs_coordinator():
+                    self._transaction_manager.lookup_coordinator_for_request(next_request_handler)
+                    break
+
+            time.sleep(self.config['retry_backoff_ms'] / 1000)
+            self._metadata.request_update()
+
+        if target_node is None:
+            self._transaction_manager.retry(next_request_handler)
+
+        return True
+
+    def _maybe_abort_batches(self, exc):
+        if self._accumulator.has_incomplete:
+            log.error("%s: Aborting producer batches due to fatal error: %s", str(self), exc)
+            self._accumulator.abort_batches(exc)
+
+    def initiate_close(self):
+        """Start closing the sender (won't complete until all data is sent)."""
+        self._running = False
+        self._accumulator.close()
+        self.wakeup()
+
+    def force_close(self):
+        """Closes the sender without sending out any pending messages."""
+        self._force_close = True
+        self.initiate_close()
+
+    def add_topic(self, topic):
+        # This is generally called from a separate thread
+        # so this needs to be a thread-safe operation
+        # we assume that checking set membership across threads
+        # is ok where self._client._topics should never
+        # remove topics for a producer instance, only add them.
+        if topic not in self._client._topics:
+            self._topics_to_add.add(topic)
+            self.wakeup()
+
+    def _maybe_wait_for_producer_id(self):
+        while not self._transaction_manager.has_producer_id():
+            try:
+                node_id = self._client.least_loaded_node()
+                if node_id is None or not self._client.await_ready(node_id):
+                    log.debug("%s, Could not find an available broker to send InitProducerIdRequest to." +
+                              " Will back off and try again.", str(self))
+                    time.sleep(self._client.least_loaded_node_refresh_ms() / 1000)
+                    continue
+                version = self._client.api_version(InitProducerIdRequest, max_version=1)
+                request = InitProducerIdRequest[version](
+                    transactional_id=self.config['transactional_id'],
+                    transaction_timeout_ms=self.config['transaction_timeout_ms'],
+                )
+                response = self._client.send_and_receive(node_id, request)
+                error_type = Errors.for_code(response.error_code)
+                if error_type is Errors.NoError:
+                    self._transaction_manager.set_producer_id_and_epoch(ProducerIdAndEpoch(response.producer_id, response.producer_epoch))
+                    break
+                elif getattr(error_type, 'retriable', False):
+                    log.debug("%s: Retriable error from InitProducerId response: %s", str(self), error_type.__name__)
+                    if getattr(error_type, 'invalid_metadata', False):
+                        self._metadata.request_update()
+                else:
+                    self._transaction_manager.transition_to_fatal_error(error_type())
+                    break
+            except Errors.KafkaConnectionError:
+                log.debug("%s: Broker %s disconnected while awaiting InitProducerId response", str(self), node_id)
+            except Errors.RequestTimedOutError:
+                log.debug("%s: InitProducerId request to node %s timed out", str(self), node_id)
+            log.debug("%s: Retry InitProducerIdRequest in %sms.", str(self), self.config['retry_backoff_ms'])
+            time.sleep(self.config['retry_backoff_ms'] / 1000)
+
+    def _failed_produce(self, batches, node_id, error):
+        log.error("%s: Error sending produce request to node %d: %s", str(self), node_id, error) # trace
+        for batch in batches:
+            self._complete_batch(batch, error, -1)
+
+    def _handle_produce_response(self, node_id, send_time, batches, response):
+        """Handle a produce response."""
+        # if we have a response, parse it
+        log.debug('%s: Parsing produce response: %r', str(self), response)
+        if response:
+            batches_by_partition = dict([(batch.topic_partition, batch)
+                                         for batch in batches])
+
+            for topic, partitions in response.topics:
+                for partition_info in partitions:
+                    if response.API_VERSION < 2:
+                        partition, error_code, offset = partition_info
+                        ts = None
+                    elif 2 <= response.API_VERSION <= 4:
+                        partition, error_code, offset, ts = partition_info
+                    elif 5 <= response.API_VERSION <= 7:
+                        partition, error_code, offset, ts, _log_start_offset = partition_info
+                    else:
+                        # Currently unused / TODO: KIP-467
+                        partition, error_code, offset, ts, _log_start_offset, _record_errors, _global_error = partition_info
+                    tp = TopicPartition(topic, partition)
+                    error = Errors.for_code(error_code)
+                    batch = batches_by_partition[tp]
+                    self._complete_batch(batch, error, offset, timestamp_ms=ts)
+
+        else:
+            # this is the acks = 0 case, just complete all requests
+            for batch in batches:
+                self._complete_batch(batch, None, -1)
+
+    def _fail_batch(self, batch, exception, base_offset=None, timestamp_ms=None):
+        exception = exception if type(exception) is not type else exception()
+        if self._transaction_manager:
+            if isinstance(exception, Errors.OutOfOrderSequenceNumberError) and \
+                    not self._transaction_manager.is_transactional() and \
+                    self._transaction_manager.has_producer_id(batch.producer_id):
+                log.error("%s: The broker received an out of order sequence number for topic-partition %s"
+                          " at offset %s. This indicates data loss on the broker, and should be investigated.",
+                          str(self), batch.topic_partition, base_offset)
+
+                # Reset the transaction state since we have hit an irrecoverable exception and cannot make any guarantees
+                # about the previously committed message. Note that this will discard the producer id and sequence
+                # numbers for all existing partitions.
+                self._transaction_manager.reset_producer_id()
+            elif isinstance(exception, (Errors.ClusterAuthorizationFailedError,
+                                        Errors.TransactionalIdAuthorizationFailedError,
+                                        Errors.ProducerFencedError,
+                                        Errors.InvalidTxnStateError)):
+                self._transaction_manager.transition_to_fatal_error(exception)
+            elif self._transaction_manager.is_transactional():
+                self._transaction_manager.transition_to_abortable_error(exception)
+
+        if self._sensors:
+            self._sensors.record_errors(batch.topic_partition.topic, batch.record_count)
+
+        if batch.done(base_offset=base_offset, timestamp_ms=timestamp_ms, exception=exception):
+            self._maybe_remove_from_inflight_batches(batch)
+            self._accumulator.deallocate(batch)
+
+    def _complete_batch(self, batch, error, base_offset, timestamp_ms=None):
+        """Complete or retry the given batch of records.
+
+        Arguments:
+            batch (ProducerBatch): The record batch
+            error (Exception): The error (or None if none)
+            base_offset (int): The base offset assigned to the records if successful
+            timestamp_ms (int, optional): The timestamp returned by the broker for this batch
+        """
+        # Standardize no-error to None
+        if error is Errors.NoError:
+            error = None
+
+        if error is not None:
+            if self._can_retry(batch, error):
+                # retry
+                log.warning("%s: Got error produce response on topic-partition %s,"
+                            " retrying (%s attempts left). Error: %s",
+                            str(self), batch.topic_partition,
+                            self.config['retries'] - batch.attempts - 1,
+                            error)
+
+                # If idempotence is enabled only retry the request if the batch matches our current producer id and epoch
+                if not self._transaction_manager or self._transaction_manager.producer_id_and_epoch.match(batch):
+                    log.debug("%s: Retrying batch to topic-partition %s. Sequence number: %s",
+                              str(self), batch.topic_partition,
+                              self._transaction_manager.sequence_number(batch.topic_partition) if self._transaction_manager else None)
+                    self._accumulator.reenqueue(batch)
+                    self._maybe_remove_from_inflight_batches(batch)
+                    if self._sensors:
+                        self._sensors.record_retries(batch.topic_partition.topic, batch.record_count)
+                else:
+                    log.warning("%s: Attempted to retry sending a batch but the producer id/epoch changed from %s/%s to %s/%s. This batch will be dropped",
+                                str(self), batch.producer_id, batch.producer_epoch,
+                                self._transaction_manager.producer_id_and_epoch.producer_id,
+                                self._transaction_manager.producer_id_and_epoch.epoch)
+                    self._fail_batch(batch, error, base_offset=base_offset, timestamp_ms=timestamp_ms)
+            else:
+                if error is Errors.TopicAuthorizationFailedError:
+                    error = error(batch.topic_partition.topic)
+
+                # tell the user the result of their request
+                self._fail_batch(batch, error, base_offset=base_offset, timestamp_ms=timestamp_ms)
+
+            if error is Errors.UnknownTopicOrPartitionError:
+                log.warning("%s: Received unknown topic or partition error in produce request on partition %s."
+                            " The topic/partition may not exist or the user may not have Describe access to it",
+                            str(self), batch.topic_partition)
+
+            if getattr(error, 'invalid_metadata', False):
+                self._metadata.request_update()
+
+        else:
+            if batch.done(base_offset=base_offset, timestamp_ms=timestamp_ms):
+                self._maybe_remove_from_inflight_batches(batch)
+                self._accumulator.deallocate(batch)
+
+            if self._transaction_manager and self._transaction_manager.producer_id_and_epoch.match(batch):
+                self._transaction_manager.increment_sequence_number(batch.topic_partition, batch.record_count)
+                log.debug("%s: Incremented sequence number for topic-partition %s to %s", str(self), batch.topic_partition,
+                          self._transaction_manager.sequence_number(batch.topic_partition))
+
+        # Unmute the completed partition.
+        if self.config['guarantee_message_order']:
+            self._accumulator.muted.remove(batch.topic_partition)
+
+    def _can_retry(self, batch, error):
+        """
+        We can retry a send if the error is transient and the number of
+        attempts taken is fewer than the maximum allowed
+        """
+        return (not batch.has_reached_delivery_timeout(self._accumulator.delivery_timeout_ms) and
+                batch.attempts < self.config['retries'] and
+                batch.final_state is None and
+                getattr(error, 'retriable', False))
+
+    def _create_produce_requests(self, collated):
+        """
+        Transfer the record batches into a list of produce requests on a
+        per-node basis.
+
+        Arguments:
+            collated: {node_id: [ProducerBatch]}
+
+        Returns:
+            dict: {node_id: ProduceRequest} (version depends on client api_versions)
+        """
+        requests = {}
+        for node_id, batches in six.iteritems(collated):
+            if batches:
+                requests[node_id] = self._produce_request(
+                    node_id, self.config['acks'],
+                    self.config['request_timeout_ms'], batches)
+        return requests
+
+    def _produce_request(self, node_id, acks, timeout, batches):
+        """Create a produce request from the given record batches.
+
+        Returns:
+            ProduceRequest (version depends on client api_versions)
+        """
+        produce_records_by_partition = collections.defaultdict(dict)
+        for batch in batches:
+            topic = batch.topic_partition.topic
+            partition = batch.topic_partition.partition
+
+            buf = batch.records.buffer()
+            produce_records_by_partition[topic][partition] = buf
+
+        version = self._client.api_version(ProduceRequest, max_version=7)
+        topic_partition_data = [
+            (topic, list(partition_info.items()))
+            for topic, partition_info in six.iteritems(produce_records_by_partition)]
+        transactional_id = self._transaction_manager.transactional_id if self._transaction_manager else None
+        if version >= 3:
+            return ProduceRequest[version](
+                transactional_id=transactional_id,
+                required_acks=acks,
+                timeout=timeout,
+                topics=topic_partition_data,
+            )
+        else:
+            if transactional_id is not None:
+                log.warning('%s: Broker does not support ProduceRequest v3+, required for transactional_id', str(self))
+            return ProduceRequest[version](
+                required_acks=acks,
+                timeout=timeout,
+                topics=topic_partition_data,
+            )
+
+    def wakeup(self):
+        """Wake up the selector associated with this send thread."""
+        self._client.wakeup()
+
+    def bootstrap_connected(self):
+        return self._client.bootstrap_connected()
+
+    def __str__(self):
+        return "<Sender client_id=%s transactional_id=%s>" % (self.config['client_id'], self.config['transactional_id'])
+
+
+class SenderMetrics(object):
+
+    def __init__(self, metrics, client, metadata):
+        self.metrics = metrics
+        self._client = client
+        self._metadata = metadata
+
+        sensor_name = 'batch-size'
+        self.batch_size_sensor = self.metrics.sensor(sensor_name)
+        self.add_metric('batch-size-avg', Avg(),
+                        sensor_name=sensor_name,
+                        description='The average number of bytes sent per partition per-request.')
+        self.add_metric('batch-size-max', Max(),
+                        sensor_name=sensor_name,
+                        description='The max number of bytes sent per partition per-request.')
+
+        sensor_name = 'compression-rate'
+        self.compression_rate_sensor = self.metrics.sensor(sensor_name)
+        self.add_metric('compression-rate-avg', Avg(),
+                        sensor_name=sensor_name,
+                        description='The average compression rate of record batches.')
+
+        sensor_name = 'queue-time'
+        self.queue_time_sensor = self.metrics.sensor(sensor_name)
+        self.add_metric('record-queue-time-avg', Avg(),
+                        sensor_name=sensor_name,
+                        description='The average time in ms record batches spent in the record accumulator.')
+        self.add_metric('record-queue-time-max', Max(),
+                        sensor_name=sensor_name,
+                        description='The maximum time in ms record batches spent in the record accumulator.')
+
+        sensor_name = 'records-per-request'
+        self.records_per_request_sensor = self.metrics.sensor(sensor_name)
+        self.add_metric('record-send-rate', Rate(),
+                        sensor_name=sensor_name,
+                        description='The average number of records sent per second.')
+        self.add_metric('records-per-request-avg', Avg(),
+                        sensor_name=sensor_name,
+                        description='The average number of records per request.')
+
+        sensor_name = 'bytes'
+        self.byte_rate_sensor = self.metrics.sensor(sensor_name)
+        self.add_metric('byte-rate', Rate(),
+                        sensor_name=sensor_name,
+                        description='The average number of bytes sent per second.')
+
+        sensor_name = 'record-retries'
+        self.retry_sensor = self.metrics.sensor(sensor_name)
+        self.add_metric('record-retry-rate', Rate(),
+                        sensor_name=sensor_name,
+                        description='The average per-second number of retried record sends')
+
+        sensor_name = 'errors'
+        self.error_sensor = self.metrics.sensor(sensor_name)
+        self.add_metric('record-error-rate', Rate(),
+                        sensor_name=sensor_name,
+                        description='The average per-second number of record sends that resulted in errors')
+
+        sensor_name = 'record-size-max'
+        self.max_record_size_sensor = self.metrics.sensor(sensor_name)
+        self.add_metric('record-size-max', Max(),
+                        sensor_name=sensor_name,
+                        description='The maximum record size across all batches')
+        self.add_metric('record-size-avg', Avg(),
+                        sensor_name=sensor_name,
+                        description='The average maximum record size per batch')
+
+        self.add_metric('requests-in-flight',
+                        AnonMeasurable(lambda *_: self._client.in_flight_request_count()),
+                        description='The current number of in-flight requests awaiting a response.')
+
+        self.add_metric('metadata-age',
+                        AnonMeasurable(lambda _, now: (now - self._metadata._last_successful_refresh_ms) / 1000),
+                        description='The age in seconds of the current producer metadata being used.')
+
+    def add_metric(self, metric_name, measurable, group_name='producer-metrics',
+                   description=None, tags=None,
+                   sensor_name=None):
+        m = self.metrics
+        metric = m.metric_name(metric_name, group_name, description, tags)
+        if sensor_name:
+            sensor = m.sensor(sensor_name)
+            sensor.add(metric, measurable)
+        else:
+            m.add_metric(metric, measurable)
+
+    def maybe_register_topic_metrics(self, topic):
+
+        def sensor_name(name):
+            return 'topic.{0}.{1}'.format(topic, name)
+
+        # if one sensor of the metrics has been registered for the topic,
+        # then all other sensors should have been registered; and vice versa
+        if not self.metrics.get_sensor(sensor_name('records-per-batch')):
+
+            self.add_metric('record-send-rate', Rate(),
+                            sensor_name=sensor_name('records-per-batch'),
+                            group_name='producer-topic-metrics.' + topic,
+                            description= 'Records sent per second for topic ' + topic)
+
+            self.add_metric('byte-rate', Rate(),
+                            sensor_name=sensor_name('bytes'),
+                            group_name='producer-topic-metrics.' + topic,
+                            description='Bytes per second for topic ' + topic)
+
+            self.add_metric('compression-rate', Avg(),
+                            sensor_name=sensor_name('compression-rate'),
+                            group_name='producer-topic-metrics.' + topic,
+                            description='Average Compression ratio for topic ' + topic)
+
+            self.add_metric('record-retry-rate', Rate(),
+                            sensor_name=sensor_name('record-retries'),
+                            group_name='producer-topic-metrics.' + topic,
+                            description='Record retries per second for topic ' + topic)
+
+            self.add_metric('record-error-rate', Rate(),
+                            sensor_name=sensor_name('record-errors'),
+                            group_name='producer-topic-metrics.' + topic,
+                            description='Record errors per second for topic ' + topic)
+
+    def update_produce_request_metrics(self, batches_map):
+        for node_batch in batches_map.values():
+            records = 0
+            total_bytes = 0
+            for batch in node_batch:
+                # register all per-topic metrics at once
+                topic = batch.topic_partition.topic
+                self.maybe_register_topic_metrics(topic)
+
+                # per-topic record send rate
+                topic_records_count = self.metrics.get_sensor(
+                    'topic.' + topic + '.records-per-batch')
+                topic_records_count.record(batch.record_count)
+
+                # per-topic bytes send rate
+                topic_byte_rate = self.metrics.get_sensor(
+                    'topic.' + topic + '.bytes')
+                topic_byte_rate.record(batch.records.size_in_bytes())
+
+                # per-topic compression rate
+                topic_compression_rate = self.metrics.get_sensor(
+                    'topic.' + topic + '.compression-rate')
+                topic_compression_rate.record(batch.records.compression_rate())
+
+                # global metrics
+                self.batch_size_sensor.record(batch.records.size_in_bytes())
+                if batch.drained:
+                    self.queue_time_sensor.record(batch.drained - batch.created)
+                self.compression_rate_sensor.record(batch.records.compression_rate())
+                self.max_record_size_sensor.record(batch.max_record_size)
+                records += batch.record_count
+                total_bytes += batch.records.size_in_bytes()
+
+            if node_batch:
+                self.records_per_request_sensor.record(records)
+                self.byte_rate_sensor.record(total_bytes)
+
+    def record_retries(self, topic, count):
+        self.retry_sensor.record(count)
+        sensor = self.metrics.get_sensor('topic.' + topic + '.record-retries')
+        if sensor:
+            sensor.record(count)
+
+    def record_errors(self, topic, count):
+        self.error_sensor.record(count)
+        sensor = self.metrics.get_sensor('topic.' + topic + '.record-errors')
+        if sensor:
+            sensor.record(count)

venv/lib/python3.11/site-packages/kafka/producer/transaction_manager.py

@@ -0,0 +1,981 @@
+from __future__ import absolute_import, division
+
+import abc 
+import collections
+import heapq
+import logging
+import threading
+
+from kafka.vendor import six
+
+try:
+    # enum in stdlib as of py3.4
+    from enum import IntEnum  # pylint: disable=import-error
+except ImportError:
+    # vendored backport module
+    from kafka.vendor.enum34 import IntEnum
+
+import kafka.errors as Errors
+from kafka.protocol.add_offsets_to_txn import AddOffsetsToTxnRequest
+from kafka.protocol.add_partitions_to_txn import AddPartitionsToTxnRequest
+from kafka.protocol.end_txn import EndTxnRequest
+from kafka.protocol.find_coordinator import FindCoordinatorRequest
+from kafka.protocol.init_producer_id import InitProducerIdRequest
+from kafka.protocol.txn_offset_commit import TxnOffsetCommitRequest
+from kafka.structs import TopicPartition
+
+
+log = logging.getLogger(__name__)
+
+
+NO_PRODUCER_ID = -1
+NO_PRODUCER_EPOCH = -1
+NO_SEQUENCE = -1
+
+
+class ProducerIdAndEpoch(object):
+    __slots__ = ('producer_id', 'epoch')
+
+    def __init__(self, producer_id, epoch):
+        self.producer_id = producer_id
+        self.epoch = epoch
+
+    @property
+    def is_valid(self):
+        return NO_PRODUCER_ID < self.producer_id
+
+    def match(self, batch):
+        return self.producer_id == batch.producer_id and self.epoch == batch.producer_epoch
+
+    def __eq__(self, other):
+        return isinstance(other, ProducerIdAndEpoch) and self.producer_id == other.producer_id and self.epoch == other.epoch
+
+    def __str__(self):
+        return "ProducerIdAndEpoch(producer_id={}, epoch={})".format(self.producer_id, self.epoch)
+
+
+class TransactionState(IntEnum):
+    UNINITIALIZED = 0
+    INITIALIZING = 1
+    READY = 2
+    IN_TRANSACTION = 3
+    COMMITTING_TRANSACTION = 4
+    ABORTING_TRANSACTION = 5
+    ABORTABLE_ERROR = 6
+    FATAL_ERROR = 7
+
+    @classmethod
+    def is_transition_valid(cls, source, target):
+        if target == cls.INITIALIZING:
+            return source == cls.UNINITIALIZED
+        elif target == cls.READY:
+            return source in (cls.INITIALIZING, cls.COMMITTING_TRANSACTION, cls.ABORTING_TRANSACTION)
+        elif target == cls.IN_TRANSACTION:
+            return source == cls.READY
+        elif target == cls.COMMITTING_TRANSACTION:
+            return source == cls.IN_TRANSACTION
+        elif target == cls.ABORTING_TRANSACTION:
+            return source in (cls.IN_TRANSACTION, cls.ABORTABLE_ERROR)
+        elif target == cls.ABORTABLE_ERROR:
+            return source in (cls.IN_TRANSACTION, cls.COMMITTING_TRANSACTION, cls.ABORTABLE_ERROR)
+        elif target == cls.UNINITIALIZED:
+            # Disallow transitions to UNITIALIZED
+            return False
+        elif target == cls.FATAL_ERROR:
+            # We can transition to FATAL_ERROR unconditionally.
+            # FATAL_ERROR is never a valid starting state for any transition. So the only option is to close the
+            # producer or do purely non transactional requests.
+            return True
+
+
+class Priority(IntEnum):
+    # We use the priority to determine the order in which requests need to be sent out. For instance, if we have
+    # a pending FindCoordinator request, that must always go first. Next, If we need a producer id, that must go second.
+    # The endTxn request must always go last.
+    FIND_COORDINATOR = 0
+    INIT_PRODUCER_ID = 1
+    ADD_PARTITIONS_OR_OFFSETS = 2
+    END_TXN = 3
+
+
+class TransactionManager(object):
+    """
+    A class which maintains state for transactions. Also keeps the state necessary to ensure idempotent production.
+    """
+    NO_INFLIGHT_REQUEST_CORRELATION_ID = -1
+    # The retry_backoff_ms is overridden to the following value if the first AddPartitions receives a
+    # CONCURRENT_TRANSACTIONS error.
+    ADD_PARTITIONS_RETRY_BACKOFF_MS = 20
+
+    def __init__(self, transactional_id=None, transaction_timeout_ms=0, retry_backoff_ms=100, api_version=(0, 11), metadata=None):
+        self._api_version = api_version
+        self._metadata = metadata
+
+        self._sequence_numbers = collections.defaultdict(lambda: 0)
+
+        self.transactional_id = transactional_id
+        self.transaction_timeout_ms = transaction_timeout_ms
+        self._transaction_coordinator = None
+        self._consumer_group_coordinator = None
+        self._new_partitions_in_transaction = set()
+        self._pending_partitions_in_transaction = set()
+        self._partitions_in_transaction = set()
+        self._pending_txn_offset_commits = dict()
+
+        self._current_state = TransactionState.UNINITIALIZED
+        self._last_error = None
+        self.producer_id_and_epoch = ProducerIdAndEpoch(NO_PRODUCER_ID, NO_PRODUCER_EPOCH)
+
+        self._transaction_started = False
+
+        self._pending_requests = [] # priority queue via heapq
+        self._pending_requests_sort_id = 0
+        self._in_flight_request_correlation_id = self.NO_INFLIGHT_REQUEST_CORRELATION_ID
+
+        # This is used by the TxnRequestHandlers to control how long to back off before a given request is retried.
+        # For instance, this value is lowered by the AddPartitionsToTxnHandler when it receives a CONCURRENT_TRANSACTIONS
+        # error for the first AddPartitionsRequest in a transaction.
+        self.retry_backoff_ms = retry_backoff_ms
+        self._lock = threading.Condition()
+
+    def initialize_transactions(self):
+        with self._lock:
+            self._ensure_transactional()
+            self._transition_to(TransactionState.INITIALIZING)
+            self.set_producer_id_and_epoch(ProducerIdAndEpoch(NO_PRODUCER_ID, NO_PRODUCER_EPOCH))
+            self._sequence_numbers.clear()
+            handler = InitProducerIdHandler(self, self.transaction_timeout_ms)
+            self._enqueue_request(handler)
+            return handler.result
+
+    def begin_transaction(self):
+        with self._lock:
+            self._ensure_transactional()
+            self._maybe_fail_with_error()
+            self._transition_to(TransactionState.IN_TRANSACTION)
+
+    def begin_commit(self):
+        with self._lock:
+            self._ensure_transactional()
+            self._maybe_fail_with_error()
+            self._transition_to(TransactionState.COMMITTING_TRANSACTION)
+            return self._begin_completing_transaction(True)
+
+    def begin_abort(self):
+        with self._lock:
+            self._ensure_transactional()
+            if self._current_state != TransactionState.ABORTABLE_ERROR:
+                self._maybe_fail_with_error()
+            self._transition_to(TransactionState.ABORTING_TRANSACTION)
+
+            # We're aborting the transaction, so there should be no need to add new partitions
+            self._new_partitions_in_transaction.clear()
+            return self._begin_completing_transaction(False)
+
+    def _begin_completing_transaction(self, committed):
+        if self._new_partitions_in_transaction:
+            self._enqueue_request(self._add_partitions_to_transaction_handler())
+        handler = EndTxnHandler(self, committed)
+        self._enqueue_request(handler)
+        return handler.result
+
+    def send_offsets_to_transaction(self, offsets, consumer_group_id):
+        with self._lock:
+            self._ensure_transactional()
+            self._maybe_fail_with_error()
+            if self._current_state != TransactionState.IN_TRANSACTION:
+                raise Errors.KafkaError("Cannot send offsets to transaction because the producer is not in an active transaction")
+
+            log.debug("Begin adding offsets %s for consumer group %s to transaction", offsets, consumer_group_id)
+            handler = AddOffsetsToTxnHandler(self, consumer_group_id, offsets)
+            self._enqueue_request(handler)
+            return handler.result
+
+    def maybe_add_partition_to_transaction(self, topic_partition):
+        with self._lock:
+            self._fail_if_not_ready_for_send()
+
+            if self.is_partition_added(topic_partition) or self.is_partition_pending_add(topic_partition):
+                return
+
+            log.debug("Begin adding new partition %s to transaction", topic_partition)
+            self._new_partitions_in_transaction.add(topic_partition)
+
+    def _fail_if_not_ready_for_send(self):
+        with self._lock:
+            if self.has_error():
+                raise Errors.KafkaError(
+                        "Cannot perform send because at least one previous transactional or"
+                        " idempotent request has failed with errors.", self._last_error)
+
+            if self.is_transactional():
+                if not self.has_producer_id():
+                    raise Errors.IllegalStateError(
+                            "Cannot perform a 'send' before completing a call to init_transactions"
+                            " when transactions are enabled.")
+
+                if self._current_state != TransactionState.IN_TRANSACTION:
+                    raise Errors.IllegalStateError("Cannot call send in state %s" % (self._current_state.name,))
+
+    def is_send_to_partition_allowed(self, tp):
+        with self._lock:
+            if self.has_fatal_error():
+                return False
+            return not self.is_transactional() or tp in self._partitions_in_transaction
+
+    def has_producer_id(self, producer_id=None):
+        if producer_id is None:
+            return self.producer_id_and_epoch.is_valid
+        else:
+            return self.producer_id_and_epoch.producer_id == producer_id
+
+    def is_transactional(self):
+        return self.transactional_id is not None
+
+    def has_partitions_to_add(self):
+        with self._lock:
+            return bool(self._new_partitions_in_transaction) or bool(self._pending_partitions_in_transaction)
+
+    def is_completing(self):
+        with self._lock:
+            return self._current_state in (
+                TransactionState.COMMITTING_TRANSACTION,
+                TransactionState.ABORTING_TRANSACTION)
+
+    @property
+    def last_error(self):
+        return self._last_error
+
+    def has_error(self):
+        with self._lock:
+            return self._current_state in (
+                TransactionState.ABORTABLE_ERROR,
+                TransactionState.FATAL_ERROR)
+
+    def is_aborting(self):
+        with self._lock:
+            return self._current_state == TransactionState.ABORTING_TRANSACTION
+
+    def transition_to_abortable_error(self, exc):
+        with self._lock:
+            if self._current_state == TransactionState.ABORTING_TRANSACTION:
+                log.debug("Skipping transition to abortable error state since the transaction is already being "
+                          " aborted. Underlying exception: %s", exc)
+                return
+            self._transition_to(TransactionState.ABORTABLE_ERROR, error=exc)
+
+    def transition_to_fatal_error(self, exc):
+        with self._lock:
+            self._transition_to(TransactionState.FATAL_ERROR, error=exc)
+
+    def is_partition_added(self, partition):
+        with self._lock:
+            return partition in self._partitions_in_transaction
+
+    def is_partition_pending_add(self, partition):
+        return partition in self._new_partitions_in_transaction or partition in self._pending_partitions_in_transaction
+
+    def has_producer_id_and_epoch(self, producer_id, producer_epoch):
+        return (
+            self.producer_id_and_epoch.producer_id == producer_id and
+            self.producer_id_and_epoch.epoch == producer_epoch
+        )
+
+    def set_producer_id_and_epoch(self, producer_id_and_epoch):
+        if not isinstance(producer_id_and_epoch, ProducerIdAndEpoch):
+            raise TypeError("ProducerAndIdEpoch type required")
+        log.info("ProducerId set to %s with epoch %s",
+                 producer_id_and_epoch.producer_id, producer_id_and_epoch.epoch)
+        self.producer_id_and_epoch = producer_id_and_epoch
+
+    def reset_producer_id(self):
+        """
+        This method is used when the producer needs to reset its internal state because of an irrecoverable exception
+        from the broker.
+
+        We need to reset the producer id and associated state when we have sent a batch to the broker, but we either get
+        a non-retriable exception or we run out of retries, or the batch expired in the producer queue after it was already
+        sent to the broker.
+
+        In all of these cases, we don't know whether batch was actually committed on the broker, and hence whether the
+        sequence number was actually updated. If we don't reset the producer state, we risk the chance that all future
+        messages will return an OutOfOrderSequenceNumberError.
+
+        Note that we can't reset the producer state for the transactional producer as this would mean bumping the epoch
+        for the same producer id. This might involve aborting the ongoing transaction during the initProducerIdRequest,
+        and the user would not have any way of knowing this happened. So for the transactional producer,
+        it's best to return the produce error to the user and let them abort the transaction and close the producer explicitly.
+        """
+        with self._lock:
+            if self.is_transactional():
+                raise Errors.IllegalStateError( 
+                    "Cannot reset producer state for a transactional producer."
+                    " You must either abort the ongoing transaction or"
+                    " reinitialize the transactional producer instead")
+            self.set_producer_id_and_epoch(ProducerIdAndEpoch(NO_PRODUCER_ID, NO_PRODUCER_EPOCH))
+            self._sequence_numbers.clear()
+
+    def sequence_number(self, tp):
+        with self._lock:
+            return self._sequence_numbers[tp]
+
+    def increment_sequence_number(self, tp, increment):
+        with self._lock:
+            if tp not in self._sequence_numbers:
+                raise Errors.IllegalStateError("Attempt to increment sequence number for a partition with no current sequence.")
+            # Sequence number wraps at java max int
+            base = self._sequence_numbers[tp]
+            if base > (2147483647 - increment):
+              self._sequence_numbers[tp] = increment - (2147483647 - base) - 1
+            else:
+                self._sequence_numbers[tp] += increment
+
+    def next_request_handler(self, has_incomplete_batches):
+        with self._lock:
+            if self._new_partitions_in_transaction:
+                self._enqueue_request(self._add_partitions_to_transaction_handler())
+
+            if not self._pending_requests:
+                return None
+
+            _, _, next_request_handler = self._pending_requests[0]
+            # Do not send the EndTxn until all batches have been flushed
+            if isinstance(next_request_handler, EndTxnHandler) and has_incomplete_batches:
+                return None
+
+            heapq.heappop(self._pending_requests)
+            if self._maybe_terminate_request_with_error(next_request_handler):
+                log.debug("Not sending transactional request %s because we are in an error state",
+                          next_request_handler.request)
+                return None
+
+            if isinstance(next_request_handler, EndTxnHandler) and not self._transaction_started:
+                next_request_handler.result.done()
+                if self._current_state != TransactionState.FATAL_ERROR:
+                    log.debug("Not sending EndTxn for completed transaction since no partitions"
+                              " or offsets were successfully added")
+                    self._complete_transaction()
+                try:
+                    _, _, next_request_handler = heapq.heappop(self._pending_requests)
+                except IndexError:
+                    next_request_handler = None
+
+            if next_request_handler:
+                log.debug("Request %s dequeued for sending", next_request_handler.request)
+
+            return next_request_handler
+
+    def retry(self, request):
+        with self._lock:
+            request.set_retry()
+            self._enqueue_request(request)
+
+    def authentication_failed(self, exc):
+        with self._lock:
+            for _, _, request in self._pending_requests:
+                request.fatal_error(exc)
+
+    def coordinator(self, coord_type):
+        if coord_type == 'group':
+            return self._consumer_group_coordinator
+        elif coord_type == 'transaction':
+            return self._transaction_coordinator
+        else:
+            raise Errors.IllegalStateError("Received an invalid coordinator type: %s" % (coord_type,))
+
+    def lookup_coordinator_for_request(self, request):
+        self._lookup_coordinator(request.coordinator_type, request.coordinator_key)
+
+    def next_in_flight_request_correlation_id(self):
+        self._in_flight_request_correlation_id += 1
+        return self._in_flight_request_correlation_id
+
+    def clear_in_flight_transactional_request_correlation_id(self):
+        self._in_flight_request_correlation_id = self.NO_INFLIGHT_REQUEST_CORRELATION_ID
+
+    def has_in_flight_transactional_request(self):
+        return self._in_flight_request_correlation_id != self.NO_INFLIGHT_REQUEST_CORRELATION_ID
+
+    def has_fatal_error(self):
+        return self._current_state == TransactionState.FATAL_ERROR
+
+    def has_abortable_error(self):
+        return self._current_state == TransactionState.ABORTABLE_ERROR
+
+    # visible for testing
+    def _test_transaction_contains_partition(self, tp):
+        with self._lock:
+            return tp in self._partitions_in_transaction
+
+    # visible for testing
+    def _test_has_pending_offset_commits(self):
+        return bool(self._pending_txn_offset_commits)
+
+    # visible for testing
+    def _test_has_ongoing_transaction(self):
+        with self._lock:
+            # transactions are considered ongoing once started until completion or a fatal error
+            return self._current_state == TransactionState.IN_TRANSACTION or self.is_completing() or self.has_abortable_error()
+
+    # visible for testing
+    def _test_is_ready(self):
+        with self._lock:
+            return self.is_transactional() and self._current_state == TransactionState.READY
+
+    def _transition_to(self, target, error=None):
+        with self._lock:
+            if not self._current_state.is_transition_valid(self._current_state, target):
+                raise Errors.KafkaError("TransactionalId %s: Invalid transition attempted from state %s to state %s" % (
+                    self.transactional_id, self._current_state.name, target.name))
+
+            if target in (TransactionState.FATAL_ERROR, TransactionState.ABORTABLE_ERROR):
+                if error is None:
+                    raise Errors.IllegalArgumentError("Cannot transition to %s with an None exception" % (target.name,))
+                self._last_error = error
+            else:
+                self._last_error = None
+
+            if self._last_error is not None:
+                log.debug("Transition from state %s to error state %s (%s)", self._current_state.name, target.name, self._last_error)
+            else:
+                log.debug("Transition from state %s to %s", self._current_state, target)
+            self._current_state = target
+
+    def _ensure_transactional(self):
+        if not self.is_transactional():
+            raise Errors.IllegalStateError("Transactional method invoked on a non-transactional producer.")
+
+    def _maybe_fail_with_error(self):
+        if self.has_error():
+            raise Errors.KafkaError("Cannot execute transactional method because we are in an error state: %s" % (self._last_error,))
+
+    def _maybe_terminate_request_with_error(self, request_handler):
+        if self.has_error():
+            if self.has_abortable_error() and isinstance(request_handler, FindCoordinatorHandler):
+                # No harm letting the FindCoordinator request go through if we're expecting to abort
+                return False
+            request_handler.fail(self._last_error)
+            return True
+        return False
+
+    def _next_pending_requests_sort_id(self):
+        self._pending_requests_sort_id += 1
+        return self._pending_requests_sort_id
+
+    def _enqueue_request(self, request_handler):
+        log.debug("Enqueuing transactional request %s", request_handler.request)
+        heapq.heappush(
+            self._pending_requests,
+            (
+                request_handler.priority, # keep lowest priority at head of queue
+                self._next_pending_requests_sort_id(), # break ties
+                request_handler
+            )
+        )
+
+    def _lookup_coordinator(self, coord_type, coord_key):
+        with self._lock:
+            if coord_type == 'group':
+                self._consumer_group_coordinator = None
+            elif coord_type == 'transaction':
+                self._transaction_coordinator = None
+            else:
+                raise Errors.IllegalStateError("Invalid coordinator type: %s" % (coord_type,))
+        self._enqueue_request(FindCoordinatorHandler(self, coord_type, coord_key))
+
+    def _complete_transaction(self):
+        with self._lock:
+            self._transition_to(TransactionState.READY)
+            self._transaction_started = False
+            self._new_partitions_in_transaction.clear()
+            self._pending_partitions_in_transaction.clear()
+            self._partitions_in_transaction.clear()
+
+    def _add_partitions_to_transaction_handler(self):
+        with self._lock:
+            self._pending_partitions_in_transaction.update(self._new_partitions_in_transaction)
+            self._new_partitions_in_transaction.clear()
+            return AddPartitionsToTxnHandler(self, self._pending_partitions_in_transaction)
+
+
+class TransactionalRequestResult(object):
+    def __init__(self):
+        self._latch = threading.Event()
+        self._error = None
+
+    def done(self, error=None):
+        self._error = error
+        self._latch.set()
+
+    def wait(self, timeout_ms=None):
+        timeout = timeout_ms / 1000 if timeout_ms is not None else None
+        success = self._latch.wait(timeout)
+        if self._error:
+            raise self._error
+        return success
+
+    @property
+    def is_done(self):
+        return self._latch.is_set()
+
+    @property
+    def succeeded(self):
+        return self._latch.is_set() and self._error is None
+
+    @property
+    def failed(self):
+        return self._latch.is_set() and self._error is not None
+
+    @property
+    def exception(self):
+        return self._error
+
+
+@six.add_metaclass(abc.ABCMeta)
+class TxnRequestHandler(object):
+    def __init__(self, transaction_manager, result=None):
+        self.transaction_manager = transaction_manager
+        self.retry_backoff_ms = transaction_manager.retry_backoff_ms
+        self.request = None
+        self._result = result or TransactionalRequestResult()
+        self._is_retry = False
+
+    @property
+    def transactional_id(self):
+        return self.transaction_manager.transactional_id
+
+    @property
+    def producer_id(self):
+        return self.transaction_manager.producer_id_and_epoch.producer_id
+
+    @property
+    def producer_epoch(self):
+        return self.transaction_manager.producer_id_and_epoch.epoch
+
+    def fatal_error(self, exc):
+        self.transaction_manager.transition_to_fatal_error(exc)
+        self._result.done(error=exc)
+
+    def abortable_error(self, exc):
+        self.transaction_manager.transition_to_abortable_error(exc)
+        self._result.done(error=exc)
+
+    def fail(self, exc):
+        self._result.done(error=exc)
+
+    def reenqueue(self):
+        with self.transaction_manager._lock:
+            self._is_retry = True
+            self.transaction_manager._enqueue_request(self)
+
+    def on_complete(self, correlation_id, response_or_exc):
+        if correlation_id != self.transaction_manager._in_flight_request_correlation_id:
+            self.fatal_error(RuntimeError("Detected more than one in-flight transactional request."))
+        else:
+            self.transaction_manager.clear_in_flight_transactional_request_correlation_id()
+            if isinstance(response_or_exc, Errors.KafkaConnectionError):
+                log.debug("Disconnected from node. Will retry.")
+                if self.needs_coordinator():
+                    self.transaction_manager._lookup_coordinator(self.coordinator_type, self.coordinator_key)
+                self.reenqueue()
+            elif isinstance(response_or_exc, Errors.UnsupportedVersionError):
+                self.fatal_error(response_or_exc)
+            elif not isinstance(response_or_exc, (Exception, type(None))):
+                log.debug("Received transactional response %s for request %s", response_or_exc, self.request)
+                with self.transaction_manager._lock:
+                    self.handle_response(response_or_exc)
+            else:
+                self.fatal_error(Errors.KafkaError("Could not execute transactional request for unknown reasons: %s" % response_or_exc))
+
+    def needs_coordinator(self):
+        return self.coordinator_type is not None
+
+    @property
+    def result(self):
+        return self._result
+
+    @property
+    def coordinator_type(self):
+        return 'transaction'
+
+    @property
+    def coordinator_key(self):
+        return self.transaction_manager.transactional_id
+
+    def set_retry(self):
+        self._is_retry = True
+
+    @property
+    def is_retry(self):
+        return self._is_retry
+
+    @abc.abstractmethod
+    def handle_response(self, response):
+        pass
+
+    @abc.abstractproperty
+    def priority(self):
+        pass
+
+
+class InitProducerIdHandler(TxnRequestHandler):
+    def __init__(self, transaction_manager, transaction_timeout_ms):
+        super(InitProducerIdHandler, self).__init__(transaction_manager)
+
+        if transaction_manager._api_version >= (2, 0):
+            version = 1
+        else:
+            version = 0
+        self.request = InitProducerIdRequest[version](
+            transactional_id=self.transactional_id,
+            transaction_timeout_ms=transaction_timeout_ms)
+
+    @property
+    def priority(self):
+        return Priority.INIT_PRODUCER_ID
+
+    def handle_response(self, response):
+        error = Errors.for_code(response.error_code)
+
+        if error is Errors.NoError:
+            self.transaction_manager.set_producer_id_and_epoch(ProducerIdAndEpoch(response.producer_id, response.producer_epoch))
+            self.transaction_manager._transition_to(TransactionState.READY)
+            self._result.done()
+        elif error in (Errors.NotCoordinatorError, Errors.CoordinatorNotAvailableError):
+            self.transaction_manager._lookup_coordinator('transaction', self.transactional_id)
+            self.reenqueue()
+        elif error in (Errors.CoordinatorLoadInProgressError, Errors.ConcurrentTransactionsError):
+            self.reenqueue()
+        elif error is Errors.TransactionalIdAuthorizationFailedError:
+            self.fatal_error(error())
+        else:
+            self.fatal_error(Errors.KafkaError("Unexpected error in InitProducerIdResponse: %s" % (error())))
+
+class AddPartitionsToTxnHandler(TxnRequestHandler):
+    def __init__(self, transaction_manager, topic_partitions):
+        super(AddPartitionsToTxnHandler, self).__init__(transaction_manager)
+
+        if transaction_manager._api_version >= (2, 7):
+            version = 2
+        elif transaction_manager._api_version >= (2, 0):
+            version = 1
+        else:
+            version = 0
+        topic_data = collections.defaultdict(list)
+        for tp in topic_partitions:
+            topic_data[tp.topic].append(tp.partition)
+        self.request = AddPartitionsToTxnRequest[version](
+            transactional_id=self.transactional_id,
+            producer_id=self.producer_id,
+            producer_epoch=self.producer_epoch,
+            topics=list(topic_data.items()))
+
+    @property
+    def priority(self):
+        return Priority.ADD_PARTITIONS_OR_OFFSETS
+
+    def handle_response(self, response):
+        has_partition_errors = False
+        unauthorized_topics = set()
+        self.retry_backoff_ms = self.transaction_manager.retry_backoff_ms
+
+        results = {TopicPartition(topic, partition): Errors.for_code(error_code)
+                   for topic, partition_data in response.results
+                   for partition, error_code in partition_data}
+
+        for tp, error in six.iteritems(results):
+            if error is Errors.NoError:
+                continue
+            elif error in (Errors.CoordinatorNotAvailableError, Errors.NotCoordinatorError):
+                self.transaction_manager._lookup_coordinator('transaction', self.transactional_id)
+                self.reenqueue()
+                return
+            elif error is Errors.ConcurrentTransactionsError:
+                self.maybe_override_retry_backoff_ms()
+                self.reenqueue()
+                return
+            elif error in (Errors.CoordinatorLoadInProgressError, Errors.UnknownTopicOrPartitionError):
+                self.reenqueue()
+                return
+            elif error is Errors.InvalidProducerEpochError:
+                self.fatal_error(error())
+                return
+            elif error is Errors.TransactionalIdAuthorizationFailedError:
+                self.fatal_error(error())
+                return
+            elif error in (Errors.InvalidProducerIdMappingError, Errors.InvalidTxnStateError):
+                self.fatal_error(Errors.KafkaError(error()))
+                return
+            elif error is Errors.TopicAuthorizationFailedError:
+                unauthorized_topics.add(tp.topic)
+            elif error is Errors.OperationNotAttemptedError:
+                log.debug("Did not attempt to add partition %s to transaction because other partitions in the"
+                          " batch had errors.", tp)
+                has_partition_errors = True
+            else:
+                log.error("Could not add partition %s due to unexpected error %s", tp, error())
+                has_partition_errors = True
+
+        partitions = set(results)
+
+        # Remove the partitions from the pending set regardless of the result. We use the presence
+        # of partitions in the pending set to know when it is not safe to send batches. However, if
+        # the partitions failed to be added and we enter an error state, we expect the batches to be
+        # aborted anyway. In this case, we must be able to continue sending the batches which are in
+        # retry for partitions that were successfully added.
+        self.transaction_manager._pending_partitions_in_transaction -= partitions
+
+        if unauthorized_topics:
+            self.abortable_error(Errors.TopicAuthorizationFailedError(unauthorized_topics))
+        elif has_partition_errors:
+            self.abortable_error(Errors.KafkaError("Could not add partitions to transaction due to errors: %s" % (results)))
+        else:
+            log.debug("Successfully added partitions %s to transaction", partitions)
+            self.transaction_manager._partitions_in_transaction.update(partitions)
+            self.transaction_manager._transaction_started = True
+            self._result.done()
+
+    def maybe_override_retry_backoff_ms(self):
+        # We only want to reduce the backoff when retrying the first AddPartition which errored out due to a
+        # CONCURRENT_TRANSACTIONS error since this means that the previous transaction is still completing and
+        # we don't want to wait too long before trying to start the new one.
+        #
+        # This is only a temporary fix, the long term solution is being tracked in
+        # https://issues.apache.org/jira/browse/KAFKA-5482
+        if not self.transaction_manager._partitions_in_transaction:
+            self.retry_backoff_ms = min(self.transaction_manager.ADD_PARTITIONS_RETRY_BACKOFF_MS, self.retry_backoff_ms)
+
+
+class FindCoordinatorHandler(TxnRequestHandler):
+    def __init__(self, transaction_manager, coord_type, coord_key):
+        super(FindCoordinatorHandler, self).__init__(transaction_manager)
+
+        self._coord_type = coord_type
+        self._coord_key = coord_key
+        if transaction_manager._api_version >= (2, 0):
+            version = 2
+        else:
+            version = 1
+        if coord_type == 'group':
+            coord_type_int8 = 0
+        elif coord_type == 'transaction':
+            coord_type_int8 = 1
+        else:
+            raise ValueError("Unrecognized coordinator type: %s" % (coord_type,))
+        self.request = FindCoordinatorRequest[version](
+            coordinator_key=coord_key,
+            coordinator_type=coord_type_int8,
+        )
+
+    @property
+    def priority(self):
+        return Priority.FIND_COORDINATOR
+
+    @property
+    def coordinator_type(self):
+        return None
+
+    @property
+    def coordinator_key(self):
+        return None
+
+    def handle_response(self, response):
+        error = Errors.for_code(response.error_code)
+
+        if error is Errors.NoError:
+            coordinator_id = self.transaction_manager._metadata.add_coordinator(
+                response, self._coord_type, self._coord_key)
+            if self._coord_type == 'group':
+                self.transaction_manager._consumer_group_coordinator = coordinator_id
+            elif self._coord_type == 'transaction':
+                self.transaction_manager._transaction_coordinator = coordinator_id
+            self._result.done()
+        elif error is Errors.CoordinatorNotAvailableError:
+            self.reenqueue()
+        elif error is Errors.TransactionalIdAuthorizationFailedError:
+            self.fatal_error(error())
+        elif error is Errors.GroupAuthorizationFailedError:
+            self.abortable_error(error(self._coord_key))
+        else:
+            self.fatal_error(Errors.KafkaError(
+                "Could not find a coordinator with type %s with key %s due to"
+                " unexpected error: %s" % (self._coord_type, self._coord_key, error())))
+
+
+class EndTxnHandler(TxnRequestHandler):
+    def __init__(self, transaction_manager, committed):
+        super(EndTxnHandler, self).__init__(transaction_manager)
+
+        if self.transaction_manager._api_version >= (2, 7):
+            version = 2
+        elif self.transaction_manager._api_version >= (2, 0):
+            version = 1
+        else:
+            version = 0
+        self.request = EndTxnRequest[version](
+            transactional_id=self.transactional_id,
+            producer_id=self.producer_id,
+            producer_epoch=self.producer_epoch,
+            committed=committed)
+
+    @property
+    def priority(self):
+        return Priority.END_TXN
+
+    def handle_response(self, response):
+        error = Errors.for_code(response.error_code)
+
+        if error is Errors.NoError:
+            self.transaction_manager._complete_transaction()
+            self._result.done()
+        elif error in (Errors.CoordinatorNotAvailableError, Errors.NotCoordinatorError):
+            self.transaction_manager._lookup_coordinator('transaction', self.transactional_id)
+            self.reenqueue()
+        elif error in (Errors.CoordinatorLoadInProgressError, Errors.ConcurrentTransactionsError):
+            self.reenqueue()
+        elif error is Errors.InvalidProducerEpochError:
+            self.fatal_error(error())
+        elif error is Errors.TransactionalIdAuthorizationFailedError:
+            self.fatal_error(error())
+        elif error is Errors.InvalidTxnStateError:
+            self.fatal_error(error())
+        else:
+            self.fatal_error(Errors.KafkaError("Unhandled error in EndTxnResponse: %s" % (error())))
+
+
+class AddOffsetsToTxnHandler(TxnRequestHandler):
+    def __init__(self, transaction_manager, consumer_group_id, offsets):
+        super(AddOffsetsToTxnHandler, self).__init__(transaction_manager)
+
+        self.consumer_group_id = consumer_group_id
+        self.offsets = offsets
+        if self.transaction_manager._api_version >= (2, 7):
+            version = 2
+        elif self.transaction_manager._api_version >= (2, 0):
+            version = 1
+        else:
+            version = 0
+        self.request = AddOffsetsToTxnRequest[version](
+            transactional_id=self.transactional_id,
+            producer_id=self.producer_id,
+            producer_epoch=self.producer_epoch,
+            group_id=consumer_group_id)
+
+    @property
+    def priority(self):
+        return Priority.ADD_PARTITIONS_OR_OFFSETS
+
+    def handle_response(self, response):
+        error = Errors.for_code(response.error_code)
+
+        if error is Errors.NoError:
+            log.debug("Successfully added partition for consumer group %s to transaction", self.consumer_group_id)
+
+            # note the result is not completed until the TxnOffsetCommit returns
+            for tp, offset in six.iteritems(self.offsets):
+                self.transaction_manager._pending_txn_offset_commits[tp] = offset
+            handler = TxnOffsetCommitHandler(self.transaction_manager, self.consumer_group_id,
+                                             self.transaction_manager._pending_txn_offset_commits, self._result)
+            self.transaction_manager._enqueue_request(handler)
+            self.transaction_manager._transaction_started = True
+        elif error in (Errors.CoordinatorNotAvailableError, Errors.NotCoordinatorError):
+            self.transaction_manager._lookup_coordinator('transaction', self.transactional_id)
+            self.reenqueue()
+        elif error in (Errors.CoordinatorLoadInProgressError, Errors.ConcurrentTransactionsError):
+            self.reenqueue()
+        elif error is Errors.InvalidProducerEpochError:
+            self.fatal_error(error())
+        elif error is Errors.TransactionalIdAuthorizationFailedError:
+            self.fatal_error(error())
+        elif error is Errors.GroupAuthorizationFailedError:
+            self.abortable_error(error(self.consumer_group_id))
+        else:
+            self.fatal_error(Errors.KafkaError("Unexpected error in AddOffsetsToTxnResponse: %s" % (error())))
+
+
+class TxnOffsetCommitHandler(TxnRequestHandler):
+    def __init__(self, transaction_manager, consumer_group_id, offsets, result):
+        super(TxnOffsetCommitHandler, self).__init__(transaction_manager, result=result)
+
+        self.consumer_group_id = consumer_group_id
+        self.offsets = offsets
+        self.request = self._build_request()
+
+    def _build_request(self):
+        if self.transaction_manager._api_version >= (2, 1):
+            version = 2
+        elif self.transaction_manager._api_version >= (2, 0):
+            version = 1
+        else:
+            version = 0
+
+        topic_data = collections.defaultdict(list)
+        for tp, offset in six.iteritems(self.offsets):
+            if version >= 2:
+                partition_data = (tp.partition, offset.offset, offset.leader_epoch, offset.metadata)
+            else:
+                partition_data = (tp.partition, offset.offset, offset.metadata)
+            topic_data[tp.topic].append(partition_data)
+
+        return TxnOffsetCommitRequest[version](
+            transactional_id=self.transactional_id,
+            group_id=self.consumer_group_id,
+            producer_id=self.producer_id,
+            producer_epoch=self.producer_epoch,
+            topics=list(topic_data.items()))
+
+    @property
+    def priority(self):
+        return Priority.ADD_PARTITIONS_OR_OFFSETS
+
+    @property
+    def coordinator_type(self):
+        return 'group'
+
+    @property
+    def coordinator_key(self):
+        return self.consumer_group_id
+
+    def handle_response(self, response):
+        lookup_coordinator = False
+        retriable_failure = False
+
+        errors = {TopicPartition(topic, partition): Errors.for_code(error_code)
+                  for topic, partition_data in response.topics
+                  for partition, error_code in partition_data}
+
+        for tp, error in six.iteritems(errors):
+            if error is Errors.NoError:
+                log.debug("Successfully added offsets for %s from consumer group %s to transaction.",
+                          tp, self.consumer_group_id)
+                del self.transaction_manager._pending_txn_offset_commits[tp]
+            elif error in (errors.CoordinatorNotAvailableError, Errors.NotCoordinatorError, Errors.RequestTimedOutError):
+                retriable_failure = True
+                lookup_coordinator = True
+            elif error is Errors.UnknownTopicOrPartitionError:
+                retriable_failure = True
+            elif error is Errors.GroupAuthorizationFailedError:
+                self.abortable_error(error(self.consumer_group_id))
+                return
+            elif error in (Errors.TransactionalIdAuthorizationFailedError,
+                           Errors.InvalidProducerEpochError,
+                           Errors.UnsupportedForMessageFormatError):
+                self.fatal_error(error())
+                return
+            else:
+                self.fatal_error(Errors.KafkaError("Unexpected error in TxnOffsetCommitResponse: %s" % (error())))
+                return
+
+        if lookup_coordinator:
+            self.transaction_manager._lookup_coordinator('group', self.consumer_group_id)
+
+        if not retriable_failure:
+            # all attempted partitions were either successful, or there was a fatal failure.
+            # either way, we are not retrying, so complete the request.
+            self.result.done()
+
+        # retry the commits which failed with a retriable error.
+        elif self.transaction_manager._pending_txn_offset_commits:
+            self.offsets = self.transaction_manager._pending_txn_offset_commits
+            self.request = self._build_request()
+            self.reenqueue()