KAFKA-12548; Propagate record error messages to application

Author: dpkp

kafka/producer/future.py

@@ -29,32 +29,35 @@ def wait(self, timeout=None):
 
 
 class FutureRecordMetadata(Future):
-    def __init__(self, produce_future, relative_offset, timestamp_ms, checksum, serialized_key_size, serialized_value_size, serialized_header_size):
+    def __init__(self, produce_future, batch_index, 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)
+        self.args = (batch_index, 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
+    def _produce_success(self, result):
+        offset, produce_timestamp_ms, record_exceptions_fn = result
 
         # Unpacking from args tuple is minor speed optimization
-        (relative_offset, timestamp_ms, checksum,
+        (batch_index, 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)
+        if record_exceptions_fn is not None:
+            self.failure(record_exceptions_fn(batch_index))
+        else:
+            # 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 batch_index is not None:
+                offset += batch_index
+            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):

kafka/producer/producer_batch.py

@@ -83,9 +83,37 @@ def abort(self, exception):
         self._final_state = FinalState.ABORTED
 
         log.debug("Aborting batch for partition %s: %s", self.topic_partition, exception)
-        self._complete_future(-1, -1, exception)
+        self._complete_future(-1, -1, lambda _: exception)
 
-    def done(self, base_offset=None, timestamp_ms=None, exception=None):
+    def complete(self, base_offset, log_append_time):
+        """Complete the batch successfully.
+
+        Arguments:
+            base_offset (int): The base offset of the messages assigned by the server
+            log_append_time (int): The log append time or -1 if CreateTime is being used
+
+        Returns: True if the batch was completed as a result of this call, and False
+            if it had been completed previously.
+        """
+        return self.done(base_offset=base_offset, timestamp_ms=log_append_time)
+
+    def complete_exceptionally(self, top_level_exception, record_exceptions_fn):
+        """
+        Complete the batch exceptionally. The provided top-level exception will be used
+        for each record future contained in the batch.
+
+        Arguments:
+            top_level_exception (Exception): top-level partition error.
+            record_exceptions_fn (callable int -> Exception): Record exception function mapping
+                batch_index to the respective record exception.
+        Returns: True if the batch was completed as a result of this call, and False
+            if it had been completed previously.
+        """
+        assert isinstance(top_level_exception, Exception)
+        assert callable(record_exceptions_fn)
+        return self.done(top_level_exception=top_level_exception, record_exceptions_fn=record_exceptions_fn)
+
+    def done(self, base_offset=None, timestamp_ms=None, top_level_exception=None, record_exceptions_fn=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
@@ -100,15 +128,15 @@ def done(self, base_offset=None, timestamp_ms=None, exception=None):
         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
+        final_state = FinalState.SUCCEEDED if top_level_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)
+                            self.topic_partition, base_offset, top_level_exception)
+            self._complete_future(base_offset, timestamp_ms, record_exceptions_fn)
             return True
 
         elif self._final_state is not FinalState.SUCCEEDED:
@@ -125,13 +153,10 @@ def done(self, base_offset=None, timestamp_ms=None, exception=None):
             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):
+    def _complete_future(self, base_offset, timestamp_ms, record_exceptions_fn):
         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)
+        self.produce_future.success((base_offset, timestamp_ms, record_exceptions_fn))
 
     def has_reached_delivery_timeout(self, delivery_timeout_ms, now=None):
         now = time.time() if now is None else now

kafka/producer/sender.py

@@ -436,6 +436,23 @@ def _handle_produce_response(self, node_id, send_time, batches, response):
             for batch in batches:
                 self._complete_batch(batch, PartitionResponse())
 
+    def _record_exceptions_fn(self, top_level_exception, record_errors, error_message):
+        """Returns a fn mapping batch_index to exception"""
+        # When no record_errors, all batches resolve to top-level exception
+        if not record_errors:
+            return lambda _: top_level_exception
+
+        record_errors_dict = dict(record_errors)
+        def record_exceptions_fn(batch_index):
+            if batch_index not in record_errors_dict:
+                return Errors.KafkaError(
+                    "Failed to append record because it was part of a batch which had one more more invalid records")
+            record_error = record_errors_dict[batch_index]
+            err_msg = record_error or error_message or top_level_exception.description
+            exc = top_level_exception.__class__ if len(record_errors) == 1 else Errors.InvalidRecordError
+            return exc(err_msg)
+        return record_exceptions_fn
+
     def _fail_batch(self, batch, partition_response):
         if partition_response.error is Errors.TopicAuthorizationFailedError:
             exception = Errors.TopicAuthorizationFailedError(batch.topic_partition.topic)
@@ -467,7 +484,8 @@ def _fail_batch(self, batch, partition_response):
         if self._sensors:
             self._sensors.record_errors(batch.topic_partition.topic, batch.record_count)
 
-        if batch.done(base_offset=partition_response.base_offset, timestamp_ms=partition_response.log_append_time, exception=exception):
+        record_exceptions_fn = self._record_exceptions_fn(exception, partition_response.record_errors, partition_response.error_message)
+        if batch.complete_exceptionally(exception, record_exceptions_fn):
             self._maybe_remove_from_inflight_batches(batch)
             self._accumulator.deallocate(batch)
 
@@ -520,7 +538,7 @@ def _complete_batch(self, batch, partition_response):
                 self._metadata.request_update()
 
         else:
-            if batch.done(base_offset=partition_response.base_offset, timestamp_ms=partition_response.log_append_time):
+            if batch.complete(partition_response.base_offset, partition_response.log_append_time):
                 self._maybe_remove_from_inflight_batches(batch)
                 self._accumulator.deallocate(batch)
 

test/test_producer_batch.py

@@ -0,0 +1,136 @@
+# pylint: skip-file
+from __future__ import absolute_import, division
+
+import pytest
+
+from kafka.errors import IllegalStateError, KafkaError
+from kafka.producer.future import FutureRecordMetadata, RecordMetadata
+from kafka.producer.producer_batch import ProducerBatch
+from kafka.record.memory_records import MemoryRecordsBuilder
+from kafka.structs import TopicPartition
+
+
+@pytest.fixture
+def tp():
+    return TopicPartition('foo', 0)
+
+
+@pytest.fixture
+def memory_records_builder():
+    return MemoryRecordsBuilder(magic=2, compression_type=0, batch_size=100000)
+
+
+@pytest.fixture
+def batch(tp, memory_records_builder):
+    return ProducerBatch(tp, memory_records_builder)
+
+
+def test_producer_batch_producer_id(tp, memory_records_builder):
+    batch = ProducerBatch(tp, memory_records_builder)
+    assert batch.producer_id == -1
+    batch.records.set_producer_state(123, 456, 789, False)
+    assert batch.producer_id == 123
+    memory_records_builder.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.complete(123, 456)
+    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, checksum=checksum,
+        serialized_key_size=3, serialized_value_size=5, serialized_header_size=-1)
+    assert future.value == expected_metadata
+
+
+def test_producer_batch_retry(batch):
+    assert not batch.in_retry()
+    batch.retry()
+    assert batch.in_retry()
+
+
+def test_batch_abort(batch):
+    future = batch.try_append(123, None, b'msg', [])
+    batch.abort(KafkaError())
+    assert future.is_done
+
+    # subsequent completion should be ignored
+    assert not batch.complete(500, 2342342341)
+    assert not batch.complete_exceptionally(KafkaError('top_level'), lambda _: KafkaError('record'))
+
+    assert future.is_done
+    with pytest.raises(KafkaError):
+        future.get()
+
+
+def test_batch_cannot_abort_twice(batch):
+    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(batch):
+    future = batch.try_append(123, None, b'msg', [])
+    batch.complete(500, 10)
+    with pytest.raises(IllegalStateError):
+        batch.complete(1000, 20)
+    record_metadata = future.get()
+    assert record_metadata.offset == 500
+    assert record_metadata.timestamp == 10
+
+
+def _test_complete_exceptionally(batch, record_count, top_level_exception, record_exceptions_fn):
+    futures = []
+    for i in range(record_count):
+        futures.append(batch.try_append(0, b'key', b'value', []))
+
+    assert record_count == batch.record_count
+
+    batch.complete_exceptionally(top_level_exception, record_exceptions_fn)
+    assert batch.is_done
+
+    for i, future in enumerate(futures):
+        assert future.is_done
+        assert future.failed()
+        assert isinstance(future.exception, RuntimeError)
+        assert record_exceptions_fn(i) == future.exception
+
+
+def test_complete_exceptionally_with_record_errors(batch):
+    record_count = 5
+    top_level_exception = RuntimeError()
+
+    record_exceptions_map = {0: RuntimeError(), 3: RuntimeError()}
+    record_exceptions_fn = lambda i: record_exceptions_map.get(i, top_level_exception)
+
+    _test_complete_exceptionally(batch, record_count, top_level_exception, record_exceptions_fn)
+
+
+def test_complete_exceptionally_with_null_record_errors(batch):
+    record_count = 5
+    top_level_exception = RuntimeError()
+
+    with pytest.raises(AssertionError):
+        _test_complete_exceptionally(batch, record_count, top_level_exception, None)