transaction_manager.py

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: ", 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.transactiona_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.TopicAuthorizationError(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(Errors.GroupAuthorizationError(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(Errors.GroupAuthorizationError(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(Errors.GroupAuthorizationError(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()