Repository: dpkp/kafka-python
Pull request #2569 (2025-03-27 04:27:55)

KIP-98: Add idempotent producer support
#1396

"Idempotence" here means that the KafkaProducer avoids causing duplicate messages on the broker via retries by using a combination of producer_id and sequence numbers.


Patch:
diff --git a/kafka/client_async.py b/kafka/client_async.py
index 19508b2..30258b7 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -27,7 +27,7 @@ from kafka.metrics.stats import Avg, Count, Rate
 from kafka.metrics.stats.rate import TimeUnit
 from kafka.protocol.broker_api_versions import BROKER_API_VERSIONS
 from kafka.protocol.metadata import MetadataRequest
-from kafka.util import Dict, WeakMethod, ensure_valid_topic_name
+from kafka.util import Dict, WeakMethod, ensure_valid_topic_name, timeout_ms_fn
 # Although this looks unused, it actually monkey-patches socket.socketpair()
 # and should be left in as long as we're using socket.socketpair() in this file
 from kafka.vendor import socketpair # noqa: F401
@@ -400,6 +400,11 @@ class KafkaClient(object):
             return True
         return False
 
+    def connection_failed(self, node_id):
+        if node_id not in self._conns:
+            return False
+        return self._conns[node_id].connect_failed()
+
     def _should_recycle_connection(self, conn):
         # Never recycle unless disconnected
         if not conn.disconnected():
@@ -1157,6 +1162,39 @@ class KafkaClient(object):
         else:
             return False
 
+    def await_ready(self, node_id, timeout_ms=30000):
+        """
+        Invokes `poll` to discard pending disconnects, followed by `client.ready` and 0 or more `client.poll`
+        invocations until the connection to `node` is ready, the timeoutMs expires or the connection fails.
+
+        It returns `true` if the call completes normally or `false` if the timeoutMs expires. If the connection fails,
+        an `IOException` is thrown instead. Note that if the `NetworkClient` has been configured with a positive
+        connection timeoutMs, it is possible for this method to raise an `IOException` for a previous connection which
+        has recently disconnected.
+
+        This method is useful for implementing blocking behaviour on top of the non-blocking `NetworkClient`, use it with
+        care.
+        """
+        inner_timeout_ms = timeout_ms_fn(timeout_ms, None)
+        self.poll(timeout_ms=0)
+        if self.is_ready(node_id):
+            return True
+
+        while not self.is_ready(node_id) and inner_timeout_ms() > 0:
+            if self.connection_failed(node_id):
+                raise Errors.KafkaConnectionError("Connection to %s failed." % (node_id,))
+            self.maybe_connect(node_id)
+            self.poll(timeout_ms=inner_timeout_ms())
+        return self.is_ready(node_id)
+
+    def send_and_receive(self, node_id, request):
+        future = self.send(node_id, request)
+        self.poll(future=future)
+        assert future.is_done
+        if future.failed():
+            raise future.exception
+        return future.value
+
 
 # OrderedDict requires python2.7+
 try:
diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py
index f0eb37a..320a165 100644
--- a/kafka/producer/kafka.py
+++ b/kafka/producer/kafka.py
@@ -19,6 +19,7 @@ from kafka.partitioner.default import DefaultPartitioner
 from kafka.producer.future import FutureRecordMetadata, FutureProduceResult
 from kafka.producer.record_accumulator import AtomicInteger, RecordAccumulator
 from kafka.producer.sender import Sender
+from kafka.producer.transaction_state import TransactionState
 from kafka.record.default_records import DefaultRecordBatchBuilder
 from kafka.record.legacy_records import LegacyRecordBatchBuilder
 from kafka.serializer import Serializer
@@ -93,6 +94,19 @@ class KafkaProducer(object):
         value_serializer (callable): used to convert user-supplied message
             values to bytes. If not None, called as f(value), should return
             bytes. Default: None.
+        enable_idempotence (bool): When set to True, the producer will ensure
+            that exactly one copy of each message is written in the stream.
+            If False, producer retries due to broker failures, etc., may write
+            duplicates of the retried message in the stream. Default: False.
+
+            Note that enabling idempotence requires
+            `max_in_flight_requests_per_connection` to be set to 1 and `retries`
+            cannot be zero. Additionally, `acks` must be set to 'all'. If these
+            values are left at their defaults, the producer will override the
+            defaults to be suitable. If the values are set to something
+            incompatible with the idempotent producer, a KafkaConfigurationError
+            will be raised.
+
         acks (0, 1, 'all'): The number of acknowledgments the producer requires
             the leader to have received before considering a request complete.
             This controls the durability of records that are sent. The
@@ -303,6 +317,7 @@ class KafkaProducer(object):
         'client_id': None,
         'key_serializer': None,
         'value_serializer': None,
+        'enable_idempotence': False,
         'acks': 1,
         'bootstrap_topics_filter': set(),
         'compression_type': None,
@@ -365,6 +380,7 @@ class KafkaProducer(object):
     def __init__(self, **configs):
         log.debug("Starting the Kafka producer")  # trace
         self.config = copy.copy(self.DEFAULT_CONFIG)
+        user_provided_configs = set(configs.keys())
         for key in self.config:
             if key in configs:
                 self.config[key] = configs.pop(key)
@@ -428,13 +444,41 @@ class KafkaProducer(object):
             assert checker(), "Libraries for {} compression codec not found".format(ct)
             self.config['compression_attrs'] = compression_attrs
 
-        message_version = self._max_usable_produce_magic()
-        self._accumulator = RecordAccumulator(message_version=message_version, **self.config)
+        self._transaction_state = None
+        if self.config['enable_idempotence']:
+            self._transaction_state = TransactionState()
+            if 'retries' not in user_provided_configs:
+                log.info("Overriding the default 'retries' config to 3 since the idempotent producer is enabled.")
+                self.config['retries'] = 3
+            elif self.config['retries'] == 0:
+                raise Errors.KafkaConfigurationError("Must set 'retries' to non-zero when using the idempotent producer.")
+
+            if 'max_in_flight_requests_per_connection' not in user_provided_configs:
+                log.info("Overriding the default 'max_in_flight_requests_per_connection' to 1 since idempontence is enabled.")
+                self.config['max_in_flight_requests_per_connection'] = 1
+            elif self.config['max_in_flight_requests_per_connection'] != 1:
+                raise Errors.KafkaConfigurationError("Must set 'max_in_flight_requests_per_connection' to 1 in order"
+                                                     " to use the idempotent producer."
+                                                     " Otherwise we cannot guarantee idempotence.")
+
+            if 'acks' not in user_provided_configs:
+                log.info("Overriding the default 'acks' config to 'all' since idempotence is enabled")
+                self.config['acks'] = -1
+            elif self.config['acks'] != -1:
+                raise Errors.KafkaConfigurationError("Must set 'acks' config to 'all' in order to use the idempotent"
+                                                     " producer. Otherwise we cannot guarantee idempotence")
+
+        message_version = self.max_usable_produce_magic(self.config['api_version'])
+        self._accumulator = RecordAccumulator(
+                transaction_state=self._transaction_state,
+                message_version=message_version,
+                **self.config)
         self._metadata = client.cluster
         guarantee_message_order = bool(self.config['max_in_flight_requests_per_connection'] == 1)
         self._sender = Sender(client, self._metadata,
                               self._accumulator,
                               metrics=self._metrics,
+                              transaction_state=self._transaction_state,
                               guarantee_message_order=guarantee_message_order,
                               **self.config)
         self._sender.daemon = True
@@ -548,16 +592,17 @@ class KafkaProducer(object):
         max_wait = self.config['max_block_ms'] / 1000
         return self._wait_on_metadata(topic, max_wait)
 
-    def _max_usable_produce_magic(self):
-        if self.config['api_version'] >= (0, 11):
+    @classmethod
+    def max_usable_produce_magic(cls, api_version):
+        if api_version >= (0, 11):
             return 2
-        elif self.config['api_version'] >= (0, 10, 0):
+        elif api_version >= (0, 10, 0):
             return 1
         else:
             return 0
 
     def _estimate_size_in_bytes(self, key, value, headers=[]):
-        magic = self._max_usable_produce_magic()
+        magic = self.max_usable_produce_magic(self.config['api_version'])
         if magic == 2:
             return DefaultRecordBatchBuilder.estimate_size_in_bytes(
                 key, value, headers)
diff --git a/kafka/producer/record_accumulator.py b/kafka/producer/record_accumulator.py
index ba82350..60fa0a3 100644
--- a/kafka/producer/record_accumulator.py
+++ b/kafka/producer/record_accumulator.py
@@ -35,9 +35,9 @@ class AtomicInteger(object):
 
 
 class ProducerBatch(object):
-    def __init__(self, tp, records):
+    def __init__(self, tp, records, now=None):
         self.max_record_size = 0
-        now = time.time()
+        now = time.time() if now is None else now
         self.created = now
         self.drained = None
         self.attempts = 0
@@ -52,13 +52,18 @@ class ProducerBatch(object):
     def record_count(self):
         return self.records.next_offset()
 
-    def try_append(self, timestamp_ms, key, value, headers):
+    @property
+    def producer_id(self):
+        return self.records.producer_id if self.records else None
+
+    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 = time.time()
+        self.last_append = now
         future = FutureRecordMetadata(self.produce_future, metadata.offset,
                                       metadata.timestamp, metadata.crc,
                                       len(key) if key is not None else -1,
@@ -81,7 +86,7 @@ class ProducerBatch(object):
                         log_start_offset, exception)  # trace
             self.produce_future.failure(exception)
 
-    def maybe_expire(self, request_timeout_ms, retry_backoff_ms, linger_ms, is_full):
+    def maybe_expire(self, request_timeout_ms, retry_backoff_ms, linger_ms, is_full, now=None):
         """Expire batches if metadata is not available
 
         A batch whose metadata is not available should be expired if one
@@ -93,7 +98,7 @@ class ProducerBatch(object):
           * the batch is in retry AND request timeout has elapsed after the
             backoff period ended.
         """
-        now = time.time()
+        now = time.time() if now is None else now
         since_append = now - self.last_append
         since_ready = now - (self.created + linger_ms / 1000.0)
         since_backoff = now - (self.last_attempt + retry_backoff_ms / 1000.0)
@@ -121,6 +126,10 @@ class ProducerBatch(object):
     def set_retry(self):
         self._retry = True
 
+    @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())
@@ -161,6 +170,7 @@ class RecordAccumulator(object):
         'compression_attrs': 0,
         'linger_ms': 0,
         'retry_backoff_ms': 100,
+        'transaction_state': None,
         'message_version': 0,
     }
 
@@ -171,6 +181,7 @@ class RecordAccumulator(object):
                 self.config[key] = configs.pop(key)
 
         self._closed = False
+        self._transaction_state = self.config['transaction_state']
         self._flushes_in_progress = AtomicInteger()
         self._appends_in_progress = AtomicInteger()
         self._batches = collections.defaultdict(collections.deque) # TopicPartition: [ProducerBatch]
@@ -233,6 +244,10 @@ class RecordAccumulator(object):
                         batch_is_full = len(dq) > 1 or last.records.is_full()
                         return future, batch_is_full, False
 
+                if self._transaction_state 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'],
@@ -310,9 +325,9 @@ class RecordAccumulator(object):
 
         return expired_batches
 
-    def reenqueue(self, batch):
+    def reenqueue(self, batch, now=None):
         """Re-enqueue the given record batch in the accumulator to retry."""
-        now = time.time()
+        now = time.time() if now is None else now
         batch.attempts += 1
         batch.last_attempt = now
         batch.last_append = now
@@ -323,7 +338,7 @@ class RecordAccumulator(object):
         with self._tp_locks[batch.topic_partition]:
             dq.appendleft(batch)
 
-    def ready(self, cluster):
+    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;
@@ -357,7 +372,7 @@ class RecordAccumulator(object):
         ready_nodes = set()
         next_ready_check = 9999999.99
         unknown_leaders_exist = False
-        now = time.time()
+        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
@@ -412,7 +427,7 @@ class RecordAccumulator(object):
                     return True
         return False
 
-    def drain(self, cluster, nodes, max_size):
+    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.
@@ -430,7 +445,7 @@ class RecordAccumulator(object):
         if not nodes:
             return {}
 
-        now = time.time()
+        now = time.time() if now is None else now
         batches = {}
         for node_id in nodes:
             size = 0
@@ -463,7 +478,26 @@ class RecordAccumulator(object):
                                     # single request
                                     break
                                 else:
+                                    producer_id_and_epoch = None
+                                    if self._transaction_state:
+                                        producer_id_and_epoch = self._transaction_state.producer_id_and_epoch
+                                        if not 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")
+                                            break
+
                                     batch = dq.popleft()
+                                    if producer_id_and_epoch 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_state.sequence_number(batch.topic_partition)
+                                        log.debug("Dest: %s: %s producer_id=%s epoch=%s sequence=%s",
+                                                  node_id, batch.topic_partition, producer_id_and_epoch.producer_id, producer_id_and_epoch.epoch,
+                                                  sequence_number)
+                                        batch.records.set_producer_state(producer_id_and_epoch.producer_id, producer_id_and_epoch.epoch, sequence_number)
                                     batch.records.close()
                                     size += batch.records.size_in_bytes()
                                     ready.append(batch)
diff --git a/kafka/producer/sender.py b/kafka/producer/sender.py
index 20af28d..24b84a9 100644
--- a/kafka/producer/sender.py
+++ b/kafka/producer/sender.py
@@ -11,6 +11,7 @@ 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.protocol.init_producer_id import InitProducerIdRequest
 from kafka.protocol.produce import ProduceRequest
 from kafka.structs import TopicPartition
 from kafka.version import __version__
@@ -29,8 +30,12 @@ class Sender(threading.Thread):
         'acks': 1,
         'retries': 0,
         'request_timeout_ms': 30000,
+        'retry_backoff_ms': 100,
         'metrics': None,
         'guarantee_message_order': False,
+        'transaction_state': None,
+        'transactional_id': None,
+        'transaction_timeout_ms': 60000,
         'client_id': 'kafka-python-' + __version__,
     }
 
@@ -52,6 +57,7 @@ class Sender(threading.Thread):
             self._sensors = SenderMetrics(self.config['metrics'], self._client, self._metadata)
         else:
             self._sensors = None
+        self._transaction_state = self.config['transaction_state']
 
     def run(self):
         """The main run loop for the sender thread."""
@@ -95,6 +101,8 @@ class Sender(threading.Thread):
         while self._topics_to_add:
             self._client.add_topic(self._topics_to_add.pop())
 
+        self._maybe_wait_for_producer_id()
+
         # get the list of partitions with data ready to send
         result = self._accumulator.ready(self._metadata)
         ready_nodes, next_ready_check_delay, unknown_leaders_exist = result
@@ -128,6 +136,13 @@ class Sender(threading.Thread):
         expired_batches = self._accumulator.abort_expired_batches(
             self.config['request_timeout_ms'], self._metadata)
 
+        # 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_state and any([batch.in_retry() for batch in expired_batches]):
+            self._transaction_state.reset_producer_id()
+            return
+
         if self._sensors:
             for expired_batch in expired_batches:
                 self._sensors.record_errors(expired_batch.topic_partition.topic, expired_batch.record_count)
@@ -185,6 +200,41 @@ class Sender(threading.Thread):
             self._topics_to_add.add(topic)
             self.wakeup()
 
+    def _maybe_wait_for_producer_id(self):
+        if not self._transaction_state:
+            return
+
+        while not self._transaction_state.has_pid():
+            try:
+                node_id = self._client.least_loaded_node()
+                if node_id is None or not self._client.await_ready(node_id):
+                    log.debug("Could not find an available broker to send InitProducerIdRequest to." +
+                              " Will back off and try again.")
+                    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_state.set_producer_id_and_epoch(response.producer_id, response.producer_epoch)
+                    return
+                elif getattr(error_type, 'retriable', False):
+                    log.debug("Retriable error from InitProducerId response: %s", error_type.__name__)
+                    if getattr(error_type, 'invalid_metadata', False):
+                        self._metadata.request_update()
+                else:
+                    log.error("Received a non-retriable error from InitProducerId response: %s", error_type.__name__)
+                    break
+            except Errors.KafkaConnectionError:
+                log.debug("Broker %s disconnected while awaiting InitProducerId response", node_id)
+            except Errors.RequestTimedOutError:
+                log.debug("InitProducerId request to node %s timed out", node_id)
+            time.sleep(self.config['retry_backoff_ms'] / 1000)
+
     def _failed_produce(self, batches, node_id, error):
         log.error("Error sending produce request to node %d: %s", node_id, error) # trace
         for batch in batches:
@@ -221,6 +271,17 @@ class Sender(threading.Thread):
             for batch in batches:
                 self._complete_batch(batch, None, -1)
 
+    def _fail_batch(batch, *args, **kwargs):
+        if self._transaction_state and self._transaction_state.producer_id_and_epoch.producer_id == batch.producer_id:
+            # 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_state.reset_producer_id()
+        batch.done(*args, **kwargs)
+        self._accumulator.deallocate(batch)
+        if self._sensors:
+            self._sensors.record_errors(batch.topic_partition.topic, batch.record_count)
+
     def _complete_batch(self, batch, error, base_offset, timestamp_ms=None, log_start_offset=None):
         """Complete or retry the given batch of records.
 
@@ -235,28 +296,55 @@ class Sender(threading.Thread):
         if error is Errors.NoError:
             error = None
 
-        if error is not None and self._can_retry(batch, error):
-            # retry
-            log.warning("Got error produce response on topic-partition %s,"
-                        " retrying (%d attempts left). Error: %s",
-                        batch.topic_partition,
-                        self.config['retries'] - batch.attempts - 1,
-                        error)
-            self._accumulator.reenqueue(batch)
-            if self._sensors:
-                self._sensors.record_retries(batch.topic_partition.topic, batch.record_count)
-        else:
-            if error is Errors.TopicAuthorizationFailedError:
-                error = error(batch.topic_partition.topic)
+        if error is not None:
+            if self._can_retry(batch, error):
+                # retry
+                log.warning("Got error produce response on topic-partition %s,"
+                            " retrying (%d attempts left). Error: %s",
+                            batch.topic_partition,
+                            self.config['retries'] - batch.attempts - 1,
+                            error)
+
+                # If idempotence is enabled only retry the request if the current PID is the same as the pid of the batch.
+                if not self._transaction_state or self._transaction_state.producer_id_and_epoch.producer_id == batch.producer_id:
+                    log.debug("Retrying batch to topic-partition %s. Sequence number: %s",
+                              batch.topic_partition,
+                              self._transaction_state.sequence_number(batch.topic_partition) if self._transaction_state else None)
+                    self._accumulator.reenqueue(batch)
+                    if self._sensors:
+                        self._sensors.record_retries(batch.topic_partition.topic, batch.record_count)
+                else:
+                    log.warning("Attempted to retry sending a batch but the producer id changed from %s to %s. This batch will be dropped" % (
+                        batch.producer_id, self._transaction_state.producer_id_and_epoch.producer_id))
+                    self._fail_batch(batch, base_offset=base_offset, timestamp_ms=timestamp_ms, exception=error, log_start_offset=log_start_offset)
+            else:
+                if error is Errors.OutOfOrderSequenceNumberError and batch.producer_id == self._transaction_state.producer_id_and_epoch.producer_id:
+                    log.error("The broker received an out of order sequence number error for produer_id %s, topic-partition %s"
+                              " at offset %s. This indicates data loss on the broker, and should be investigated.",
+                              batch.producer_id, batch.topic_partition, base_offset)
+
+                if error is Errors.TopicAuthorizationFailedError:
+                    error = error(batch.topic_partition.topic)
+
+                # tell the user the result of their request
+                self._fail_batch(batch, base_offset=base_offset, timestamp_ms=timestamp_ms, exception=error, log_start_offset=log_start_offset)
+
+            if error is Errors.UnknownTopicOrPartitionError:
+                log.warning("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",
+                            batch.topic_partition)
+
+            if getattr(error, 'invalid_metadata', False):
+                self._metadata.request_update()
 
-            # tell the user the result of their request
-            batch.done(base_offset, timestamp_ms, error, log_start_offset)
+        else:
+            batch.done(base_offset=base_offset, timestamp_ms=timestamp_ms, log_start_offset=log_start_offset)
             self._accumulator.deallocate(batch)
-            if error is not None and self._sensors:
-                self._sensors.record_errors(batch.topic_partition.topic, batch.record_count)
 
-        if getattr(error, 'invalid_metadata', False):
-            self._metadata.request_update()
+            if self._transaction_state and self._transaction_state.producer_id_and_epoch.producer_id == batch.producer_id:
+                self._transaction_state.increment_sequence_number(batch.topic_partition, batch.record_count)
+                log.debug("Incremented sequence number for topic-partition %s to %s", batch.topic_partition,
+                          self._transaction_state.sequence_number(batch.topic_partition))
 
         # Unmute the completed partition.
         if self.config['guarantee_message_order']:
diff --git a/kafka/producer/transaction_state.py b/kafka/producer/transaction_state.py
new file mode 100644
index 0000000..05cdc57
--- /dev/null
+++ b/kafka/producer/transaction_state.py
@@ -0,0 +1,96 @@
+from __future__ import absolute_import, division
+
+import collections
+import threading
+import time
+
+from kafka.errors import IllegalStateError
+
+
+NO_PRODUCER_ID = -1
+NO_PRODUCER_EPOCH = -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 __str__(self):
+        return "ProducerIdAndEpoch(producer_id={}, epoch={})".format(self.producer_id, self.epoch)
+
+class TransactionState(object):
+    __slots__ = ('producer_id_and_epoch', '_sequence_numbers', '_lock')
+
+    def __init__(self):
+        self.producer_id_and_epoch = ProducerIdAndEpoch(NO_PRODUCER_ID, NO_PRODUCER_EPOCH)
+        self._sequence_numbers = collections.defaultdict(lambda: 0)
+        self._lock = threading.Condition()
+
+    def has_pid(self):
+        return self.producer_id_and_epoch.is_valid
+
+
+    def await_producer_id_and_epoch(self, max_wait_time_ms):
+        """
+        A blocking call to get the pid and epoch for the producer. If the PID and epoch has not been set, this method
+        will block for at most maxWaitTimeMs. It is expected that this method be called from application thread
+        contexts (ie. through Producer.send). The PID it self will be retrieved in the background thread.
+
+        Arguments:
+            max_wait_time_ms (numeric): The maximum time to block.
+
+        Returns:
+            ProducerIdAndEpoch object. Callers must check the 'is_valid' property of the returned object to ensure that a
+                valid pid and epoch is actually returned.
+        """
+        with self._lock:
+            start = time.time()
+            elapsed = 0
+            while not self.has_pid() and elapsed < max_wait_time_ms:
+                self._lock.wait(max_wait_time_ms / 1000)
+                elapsed = time.time() - start
+            return self.producer_id_and_epoch
+
+    def set_producer_id_and_epoch(self, producer_id, epoch):
+        """
+        Set the pid and epoch atomically. This method will signal any callers blocked on the `pidAndEpoch` method
+        once the pid is set. This method will be called on the background thread when the broker responds with the pid.
+        """
+        with self._lock:
+            self.producer_id_and_epoch = ProducerIdAndEpoch(producer_id, epoch)
+            if self.producer_id_and_epoch.is_valid:
+                self._lock.notify_all()
+
+    def reset_producer_id(self):
+        """
+        This method is used when the producer needs to reset it's 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 OutOfOrderSequenceException.
+        """
+        with self._lock:
+            self.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 IllegalStateError("Attempt to increment sequence number for a partition with no current sequence.")
+            self._sequence_numbers[tp] += increment
diff --git a/kafka/protocol/init_producer_id.py b/kafka/protocol/init_producer_id.py
new file mode 100644
index 0000000..8426fe0
--- /dev/null
+++ b/kafka/protocol/init_producer_id.py
@@ -0,0 +1,46 @@
+from __future__ import absolute_import
+
+from kafka.protocol.api import Request, Response
+from kafka.protocol.types import Int16, Int32, Int64, Schema, String
+
+
+class InitProducerIdResponse_v0(Response):
+    API_KEY = 22
+    API_VERSION = 0
+    SCHEMA = Schema(
+        ('throttle_time_ms', Int32),
+        ('error_code', Int16),
+        ('producer_id', Int64),
+        ('producer_epoch', Int16),
+    )
+
+
+class InitProducerIdResponse_v1(Response):
+    API_KEY = 22
+    API_VERSION = 1
+    SCHEMA = InitProducerIdResponse_v0.SCHEMA
+
+
+class InitProducerIdRequest_v0(Request):
+    API_KEY = 22
+    API_VERSION = 0
+    RESPONSE_TYPE = InitProducerIdResponse_v0
+    SCHEMA = Schema(
+        ('transactional_id', String('utf-8')),
+        ('transaction_timeout_ms', Int32),
+    )
+
+
+class InitProducerIdRequest_v1(Request):
+    API_KEY = 22
+    API_VERSION = 1
+    RESPONSE_TYPE = InitProducerIdResponse_v1
+    SCHEMA = InitProducerIdRequest_v0.SCHEMA
+
+
+InitProducerIdRequest = [
+    InitProducerIdRequest_v0, InitProducerIdRequest_v1,
+]
+InitProducerIdResponse = [
+    InitProducerIdResponse_v0, InitProducerIdResponse_v1,
+]
diff --git a/kafka/record/default_records.py b/kafka/record/default_records.py
index 0d69d72..855306b 100644
--- a/kafka/record/default_records.py
+++ b/kafka/record/default_records.py
@@ -448,6 +448,15 @@ class DefaultRecordBatchBuilder(DefaultRecordBase, ABCRecordBatchBuilder):
 
         self._buffer = bytearray(self.HEADER_STRUCT.size)
 
+    def set_producer_state(self, producer_id, producer_epoch, base_sequence):
+        self._producer_id = producer_id
+        self._producer_epoch = producer_epoch
+        self._base_sequence = base_sequence
+
+    @property
+    def producer_id(self):
+        return self._producer_id
+
     def _get_attributes(self, include_compression_type=True):
         attrs = 0
         if include_compression_type:
diff --git a/kafka/record/memory_records.py b/kafka/record/memory_records.py
index 72baea5..a803047 100644
--- a/kafka/record/memory_records.py
+++ b/kafka/record/memory_records.py
@@ -22,7 +22,7 @@ from __future__ import division
 
 import struct
 
-from kafka.errors import CorruptRecordException
+from kafka.errors import CorruptRecordException, IllegalStateError, UnsupportedVersionError
 from kafka.record.abc import ABCRecords
 from kafka.record.legacy_records import LegacyRecordBatch, LegacyRecordBatchBuilder
 from kafka.record.default_records import DefaultRecordBatch, DefaultRecordBatchBuilder
@@ -113,7 +113,7 @@ class MemoryRecords(ABCRecords):
 class MemoryRecordsBuilder(object):
 
     __slots__ = ("_builder", "_batch_size", "_buffer", "_next_offset", "_closed",
-                 "_bytes_written")
+                 "_magic", "_bytes_written", "_producer_id")
 
     def __init__(self, magic, compression_type, batch_size, offset=0):
         assert magic in [0, 1, 2], "Not supported magic"
@@ -123,15 +123,18 @@ class MemoryRecordsBuilder(object):
                 magic=magic, compression_type=compression_type,
                 is_transactional=False, producer_id=-1, producer_epoch=-1,
                 base_sequence=-1, batch_size=batch_size)
+            self._producer_id = -1
         else:
             self._builder = LegacyRecordBatchBuilder(
                 magic=magic, compression_type=compression_type,
                 batch_size=batch_size)
+            self._producer_id = None
         self._batch_size = batch_size
         self._buffer = None
 
         self._next_offset = offset
         self._closed = False
+        self._magic = magic
         self._bytes_written = 0
 
     def skip(self, offsets_to_skip):
@@ -155,6 +158,24 @@ class MemoryRecordsBuilder(object):
         self._next_offset += 1
         return metadata
 
+    def set_producer_state(self, producer_id, producer_epoch, base_sequence):
+        if self._magic < 2:
+            raise UnsupportedVersionError('Producer State requires Message format v2+')
+        elif self._closed:
+            # Sequence numbers are assigned when the batch is closed while the accumulator is being drained.
+            # If the resulting ProduceRequest to the partition leader failed for a retriable error, the batch will
+            # be re queued. In this case, we should not attempt to set the state again, since changing the pid and sequence
+            # once a batch has been sent to the broker risks introducing duplicates.
+            raise IllegalStateError("Trying to set producer state of an already closed batch. This indicates a bug on the client.")
+        self._builder.set_producer_state(producer_id, producer_epoch, base_sequence)
+        self._producer_id = producer_id
+
+    @property
+    def producer_id(self):
+        if self._magic < 2:
+            raise UnsupportedVersionError('Producer State requires Message format v2+')
+        return self._producer_id
+
     def close(self):
         # This method may be called multiple times on the same batch
         # i.e., on retries
@@ -164,6 +185,8 @@ class MemoryRecordsBuilder(object):
         if not self._closed:
             self._bytes_written = self._builder.size()
             self._buffer = bytes(self._builder.build())
+            if self._magic == 2:
+                self._producer_id = self._builder.producer_id
             self._builder = None
         self._closed = True
 


Test patch:
diff --git a/test/test_producer.py b/test/test_producer.py
index 069362f..303832b 100644
--- a/test/test_producer.py
+++ b/test/test_producer.py
@@ -100,7 +100,7 @@ def test_kafka_producer_proper_record_metadata(kafka_broker, compression):
                           retries=5,
                           max_block_ms=30000,
                           compression_type=compression) as producer:
-        magic = producer._max_usable_produce_magic()
+        magic = producer.max_usable_produce_magic(producer.config['api_version'])
 
         # record headers are supported in 0.11.0
         if env_kafka_version() < (0, 11, 0):
diff --git a/test/test_record_accumulator.py b/test/test_record_accumulator.py
new file mode 100644
index 0000000..babff56
--- /dev/null
+++ b/test/test_record_accumulator.py
@@ -0,0 +1,75 @@
+# pylint: skip-file
+from __future__ import absolute_import
+
+import pytest
+import io
+
+from kafka.errors import KafkaTimeoutError
+from kafka.producer.future import FutureRecordMetadata, RecordMetadata
+from kafka.producer.record_accumulator import RecordAccumulator, ProducerBatch
+from kafka.record.memory_records import MemoryRecordsBuilder
+from kafka.structs import TopicPartition
+
+
+def test_producer_batch_producer_id():
+    tp = TopicPartition('foo', 0)
+    records = MemoryRecordsBuilder(
+        magic=2, compression_type=0, batch_size=100000)
+    batch = ProducerBatch(tp, records)
+    assert batch.producer_id == -1
+    batch.records.set_producer_state(123, 456, 789)
+    assert batch.producer_id == 123
+    records.close()
+    assert batch.producer_id == 123
+
+@pytest.mark.parametrize("magic", [0, 1, 2])
+def test_producer_batch_try_append(magic):
+    tp = TopicPartition('foo', 0)
+    records = MemoryRecordsBuilder(
+        magic=magic, compression_type=0, batch_size=100000)
+    batch = ProducerBatch(tp, records)
+    assert batch.record_count == 0
+    future = batch.try_append(0, b'key', b'value', [])
+    assert isinstance(future, FutureRecordMetadata)
+    assert not future.is_done
+    batch.done(base_offset=123, timestamp_ms=456, log_start_offset=0)
+    assert future.is_done
+    # record-level checksum only provided in v0/v1 formats; payload includes magic-byte
+    if magic == 0:
+        checksum = 592888119
+    elif magic == 1:
+        checksum = 213653215
+    else:
+        checksum = None
+
+    expected_metadata = RecordMetadata(
+        topic=tp[0], partition=tp[1], topic_partition=tp,
+        offset=123, timestamp=456, log_start_offset=0,
+        checksum=checksum,
+        serialized_key_size=3, serialized_value_size=5, serialized_header_size=-1)
+    assert future.value == expected_metadata
+
+def test_producer_batch_retry():
+    tp = TopicPartition('foo', 0)
+    records = MemoryRecordsBuilder(
+        magic=2, compression_type=0, batch_size=100000)
+    batch = ProducerBatch(tp, records)
+    assert not batch.in_retry()
+    batch.set_retry()
+    assert batch.in_retry()
+
+def test_producer_batch_maybe_expire():
+    tp = TopicPartition('foo', 0)
+    records = MemoryRecordsBuilder(
+        magic=2, compression_type=0, batch_size=100000)
+    batch = ProducerBatch(tp, records, now=1)
+    future = batch.try_append(0, b'key', b'value', [], now=2)
+    request_timeout_ms = 5000
+    retry_backoff_ms = 200
+    linger_ms = 1000
+    is_full = True
+    batch.maybe_expire(request_timeout_ms, retry_backoff_ms, linger_ms, is_full, now=20)
+    assert batch.is_done
+    assert future.is_done
+    assert future.failed()
+    assert isinstance(future.exception, KafkaTimeoutError)
diff --git a/test/test_sender.py b/test/test_sender.py
index b037d2b..eedc43d 100644
--- a/test/test_sender.py
+++ b/test/test_sender.py
@@ -6,6 +6,7 @@ import io
 
 from kafka.client_async import KafkaClient
 from kafka.protocol.broker_api_versions import BROKER_API_VERSIONS
+from kafka.producer.kafka import KafkaProducer
 from kafka.protocol.produce import ProduceRequest
 from kafka.producer.record_accumulator import RecordAccumulator, ProducerBatch
 from kafka.producer.sender import Sender
@@ -24,6 +25,7 @@ def sender(client, accumulator, metrics, mocker):
 
 
 @pytest.mark.parametrize(("api_version", "produce_version"), [
+    ((2, 1), 7),
     ((0, 10, 0), 2),
     ((0, 9), 1),
     ((0, 8, 0), 0)
@@ -31,6 +33,7 @@ def sender(client, accumulator, metrics, mocker):
 def test_produce_request(sender, mocker, api_version, produce_version):
     sender._client._api_versions = BROKER_API_VERSIONS[api_version]
     tp = TopicPartition('foo', 0)
+    magic = KafkaProducer.max_usable_produce_magic(api_version)
     records = MemoryRecordsBuilder(
         magic=1, compression_type=0, batch_size=100000)
     batch = ProducerBatch(tp, records)