KAFKA-5429 - Ignore produce response if batch was previously aborted

Author: dpkp

kafka/producer/record_accumulator.py

@@ -6,6 +6,13 @@
 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
@@ -34,6 +41,12 @@ 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):
         self.max_record_size = 0
@@ -47,6 +60,7 @@ def __init__(self, tp, records, now=None):
         self.topic_partition = tp
         self.produce_future = FutureProduceResult(tp)
         self._retry = False
+        self._final_state = None
 
     @property
     def record_count(self):
@@ -79,10 +93,29 @@ def try_append(self, timestamp_ms, key, value, headers, now=None):
                                       sum(len(h_key.encode("utf-8")) + len(h_val) for h_key, h_val in headers) if headers else -1)
         return future
 
-    def done(self, base_offset=None, timestamp_ms=None, exception=None, log_start_offset=None):
-        if self.produce_future.is_done:
-            log.warning('Batch is already closed -- ignoring batch.done()')
+    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):
+        if self._final_state is None:
+            self._final_state = FinalState.SUCCEEDED if exception is None else FinalState.FAILED
+        elif self._final_state is FinalState.ABORTED:
+            log.debug("ProduceResponse returned for %s after batch had already been aborted.", self.topic_partition)
             return
+        else:
+            raise Errors.IllegalStateError("Batch has already been completed in final state %s" % self._final_state)
+
+        self._complete_future(base_offset, timestamp_ms, exception)
+
+    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:
             log.debug("Produced messages to topic-partition %s with base offset %s", self.topic_partition, base_offset)
             self.produce_future.success((base_offset, timestamp_ms))
@@ -588,7 +621,7 @@ def _abort_batches(self, error):
             with self._tp_locks[tp]:
                 batch.records.close()
                 self._batches[tp].remove(batch)
-            batch.done(exception=error)
+            batch.abort(error)
             self.deallocate(batch)
 
     def abort_undrained_batches(self, error):
@@ -601,7 +634,7 @@ def abort_undrained_batches(self, error):
                     batch.records.close()
                     self._batches[tp].remove(batch)
             if aborted:
-                batch.done(exception=error)
+                batch.abort(error)
                 self.deallocate(batch)
 
     def close(self):

test/test_record_accumulator.py

@@ -4,7 +4,7 @@
 import pytest
 import io
 
-from kafka.errors import KafkaTimeoutError
+from kafka.errors import IllegalStateError, KafkaError, KafkaTimeoutError
 from kafka.producer.future import FutureRecordMetadata, RecordMetadata
 from kafka.producer.record_accumulator import RecordAccumulator, ProducerBatch
 from kafka.record.memory_records import MemoryRecordsBuilder
@@ -72,3 +72,54 @@ def test_producer_batch_maybe_expire():
     assert future.is_done
     assert future.failed()
     assert isinstance(future.exception, KafkaTimeoutError)
+
+def test_batch_abort():
+    tp = TopicPartition('foo', 0)
+    records = MemoryRecordsBuilder(
+        magic=2, compression_type=0, batch_size=100000)
+    batch = ProducerBatch(tp, records)
+    future = batch.try_append(123, None, b'msg', [])
+
+    batch.abort(KafkaError())
+    assert future.is_done
+
+    # subsequent completion should be ignored
+    batch.done(500, 2342342341)
+    batch.done(exception=KafkaError())
+
+    assert future.is_done
+    with pytest.raises(KafkaError):
+        future.get()
+
+def test_batch_cannot_abort_twice():
+    tp = TopicPartition('foo', 0)
+    records = MemoryRecordsBuilder(
+        magic=2, compression_type=0, batch_size=100000)
+    batch = ProducerBatch(tp, records)
+    future = batch.try_append(123, None, b'msg', [])
+
+    batch.abort(KafkaError())
+
+    with pytest.raises(IllegalStateError):
+        batch.abort(KafkaError())
+
+    assert future.is_done
+    with pytest.raises(KafkaError):
+        future.get()
+
+def test_batch_cannot_complete_twice():
+    tp = TopicPartition('foo', 0)
+    records = MemoryRecordsBuilder(
+        magic=2, compression_type=0, batch_size=100000)
+    batch = ProducerBatch(tp, records)
+    future = batch.try_append(123, None, b'msg', [])
+
+    batch.done(500, 10, None)
+
+    with pytest.raises(IllegalStateError):
+        batch.done(1000, 20, None)
+
+    record_metadata = future.get()
+
+    assert record_metadata.offset == 500
+    assert record_metadata.timestamp == 10