KAFKA-8830 make Record Headers available in onAcknowledgement (#17099)

Two sets of tests are added:
1. KafkaProducerTest
- when send success, both record.headers() and onAcknowledgement headers
are read only
- when send failure, record.headers() is writable as before and
onAcknowledgement headers is read only
2. ProducerInterceptorsTest
- make both old and new onAcknowledgement method are called successfully

Reviewers: Lianet Magrans <[email protected]>, Omnia Ibrahim
<[email protected]>, Matthias J. Sax <[email protected]>,
Andrew Schofield <[email protected]>, Chia-Ping Tsai
<[email protected]>

Author: rich-c-shop

clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java

@@ -1546,6 +1546,7 @@ private class AppendCallbacks implements RecordAccumulator.AppendCallbacks {
         private final String recordLogString;
         private volatile int partition = RecordMetadata.UNKNOWN_PARTITION;
         private volatile TopicPartition topicPartition;
+        private final Headers headers;
 
         private AppendCallbacks(Callback userCallback, ProducerInterceptors<K, V> interceptors, ProducerRecord<K, V> record) {
             this.userCallback = userCallback;
@@ -1554,6 +1555,12 @@ private AppendCallbacks(Callback userCallback, ProducerInterceptors<K, V> interc
             // whole lifetime of the batch.
             // We don't want to have an NPE here, because the interceptors would not be notified (see .doSend).
             topic = record != null ? record.topic() : null;
+            if (record != null) {
+                headers = record.headers();
+            } else {
+                headers = new RecordHeaders();
+                ((RecordHeaders) headers).setReadOnly();
+            }
             recordPartition = record != null ? record.partition() : null;
             recordLogString = log.isTraceEnabled() && record != null ? record.toString() : "";
         }
@@ -1563,7 +1570,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
             if (metadata == null) {
                 metadata = new RecordMetadata(topicPartition(), -1, -1, RecordBatch.NO_TIMESTAMP, -1, -1);
             }
-            this.interceptors.onAcknowledgement(metadata, exception);
+            this.interceptors.onAcknowledgement(metadata, exception, headers);
             if (this.userCallback != null)
                 this.userCallback.onCompletion(metadata, exception);
         }

clients/src/main/java/org/apache/kafka/clients/producer/ProducerInterceptor.java

@@ -17,6 +17,7 @@
 package org.apache.kafka.clients.producer;
 
 import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.header.Headers;
 
 /**
  * A plugin interface that allows you to intercept (and possibly mutate) the records received by the producer before
@@ -83,12 +84,37 @@ public interface ProducerInterceptor<K, V> extends Configurable, AutoCloseable {
      * @param metadata The metadata for the record that was sent (i.e. the partition and offset).
      *                 If an error occurred, metadata will contain only valid topic and maybe
      *                 partition. If partition is not given in ProducerRecord and an error occurs
-     *                 before partition gets assigned, then partition will be set to RecordMetadata.NO_PARTITION.
+     *                 before partition gets assigned, then partition will be set to {@link RecordMetadata#UNKNOWN_PARTITION}.
      *                 The metadata may be null if the client passed null record to
      *                 {@link org.apache.kafka.clients.producer.KafkaProducer#send(ProducerRecord)}.
      * @param exception The exception thrown during processing of this record. Null if no error occurred.
      */
-    void onAcknowledgement(RecordMetadata metadata, Exception exception);
+    default void onAcknowledgement(RecordMetadata metadata, Exception exception) {}
+
+    /**
+     * This method is called when the record sent to the server has been acknowledged, or when sending the record fails before
+     * it gets sent to the server.
+     * <p>
+     * This method is generally called just before the user callback is called, and in additional cases when <code>KafkaProducer.send()</code>
+     * throws an exception.
+     * <p>
+     * Any exception thrown by this method will be ignored by the caller.
+     * <p>
+     * This method will generally execute in the background I/O thread, so the implementation should be reasonably fast.
+     * Otherwise, sending of messages from other threads could be delayed.
+     *
+     * @param metadata The metadata for the record that was sent (i.e. the partition and offset).
+     *                 If an error occurred, metadata will contain only valid topic and maybe
+     *                 partition. If partition is not given in ProducerRecord and an error occurs
+     *                 before partition gets assigned, then partition will be set to {@link RecordMetadata#UNKNOWN_PARTITION}.
+     *                 The metadata may be null if the client passed null record to
+     *                 {@link org.apache.kafka.clients.producer.KafkaProducer#send(ProducerRecord)}.
+     * @param exception The exception thrown during processing of this record. Null if no error occurred.
+     * @param headers The headers for the record that was sent. It is read-only.
+     */
+    default void onAcknowledgement(RecordMetadata metadata, Exception exception, Headers headers) {
+        onAcknowledgement(metadata, exception);
+    }
 
     /**
      * This is called when interceptor is closed

clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerInterceptors.java

@@ -22,6 +22,8 @@
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.internals.Plugin;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.record.RecordBatch;
@@ -77,19 +79,20 @@ public ProducerRecord<K, V> onSend(ProducerRecord<K, V> record) {
 
     /**
      * This method is called when the record sent to the server has been acknowledged, or when sending the record fails before
-     * it gets sent to the server. This method calls {@link ProducerInterceptor#onAcknowledgement(RecordMetadata, Exception)}
+     * it gets sent to the server. This method calls {@link ProducerInterceptor#onAcknowledgement(RecordMetadata, Exception, Headers)}
      * method for each interceptor.
      *
      * This method does not throw exceptions. Exceptions thrown by any of interceptor methods are caught and ignored.
      *
      * @param metadata The metadata for the record that was sent (i.e. the partition and offset).
      *                 If an error occurred, metadata will only contain valid topic and maybe partition.
      * @param exception The exception thrown during processing of this record. Null if no error occurred.
+     * @param headers The headers for the record that was sent
      */
-    public void onAcknowledgement(RecordMetadata metadata, Exception exception) {
+    public void onAcknowledgement(RecordMetadata metadata, Exception exception, Headers headers) {
         for (Plugin<ProducerInterceptor<K, V>> interceptorPlugin : this.interceptorPlugins) {
             try {
-                interceptorPlugin.get().onAcknowledgement(metadata, exception);
+                interceptorPlugin.get().onAcknowledgement(metadata, exception, headers);
             } catch (Exception e) {
                 // do not propagate interceptor exceptions, just log
                 log.warn("Error executing interceptor onAcknowledgement callback", e);
@@ -99,7 +102,7 @@ public void onAcknowledgement(RecordMetadata metadata, Exception exception) {
 
     /**
      * This method is called when sending the record fails in {@link ProducerInterceptor#onSend
-     * (ProducerRecord)} method. This method calls {@link ProducerInterceptor#onAcknowledgement(RecordMetadata, Exception)}
+     * (ProducerRecord)} method. This method calls {@link ProducerInterceptor#onAcknowledgement(RecordMetadata, Exception, Headers)}
      * method for each interceptor
      *
      * @param record The record from client
@@ -110,14 +113,22 @@ public void onAcknowledgement(RecordMetadata metadata, Exception exception) {
     public void onSendError(ProducerRecord<K, V> record, TopicPartition interceptTopicPartition, Exception exception) {
         for (Plugin<ProducerInterceptor<K, V>> interceptorPlugin : this.interceptorPlugins) {
             try {
+                Headers headers = record != null ? record.headers() : new RecordHeaders();
+                if (headers instanceof RecordHeaders && !((RecordHeaders) headers).isReadOnly()) {
+                    // make a copy of the headers to make sure we don't change the state of origin record's headers.
+                    // original headers are still writable because client might want to mutate them before retrying.
+                    RecordHeaders recordHeaders = (RecordHeaders) headers;
+                    headers = new RecordHeaders(recordHeaders);
+                    ((RecordHeaders) headers).setReadOnly();
+                }
                 if (record == null && interceptTopicPartition == null) {
-                    interceptorPlugin.get().onAcknowledgement(null, exception);
+                    interceptorPlugin.get().onAcknowledgement(null, exception, headers);
                 } else {
                     if (interceptTopicPartition == null) {
                         interceptTopicPartition = extractTopicPartition(record);
                     }
                     interceptorPlugin.get().onAcknowledgement(new RecordMetadata(interceptTopicPartition, -1, -1,
-                                    RecordBatch.NO_TIMESTAMP, -1, -1), exception);
+                                    RecordBatch.NO_TIMESTAMP, -1, -1), exception, headers);
                 }
             } catch (Exception e) {
                 // do not propagate interceptor exceptions, just log

clients/src/main/java/org/apache/kafka/common/header/internals/RecordHeaders.java

@@ -108,6 +108,10 @@ public void setReadOnly() {
         this.isReadOnly = true;
     }
 
+    public boolean isReadOnly() {
+        return isReadOnly;
+    }
+
     public Header[] toArray() {
         return headers.isEmpty() ? Record.EMPTY_HEADERS : headers.toArray(new Header[0]);     
     }

clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java

@@ -48,7 +48,9 @@
 import org.apache.kafka.common.errors.RecordTooLargeException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
 import org.apache.kafka.common.message.ApiVersionsResponseData;
@@ -1084,13 +1086,14 @@ public void testTopicExpiryInMetadata() throws InterruptedException {
 
     @SuppressWarnings("unchecked")
     @Test
-    public void testHeaders() {
+    public void testHeadersSuccess() {
         doTestHeaders(Serializer.class);
     }
 
     private <T extends Serializer<String>> void doTestHeaders(Class<T> serializerClassToMock) {
         Map<String, Object> configs = new HashMap<>();
         configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        configs.put(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, ProducerInterceptorForHeaders.class.getName());
         Serializer<String> keySerializer = mock(serializerClassToMock);
         Serializer<String> valueSerializer = mock(serializerClassToMock);
 
@@ -1119,7 +1122,9 @@ private <T extends Serializer<String>> void doTestHeaders(Class<T> serializerCla
         producer.send(record, null);
 
         //ensure headers are closed and cannot be mutated post send
-        assertThrows(IllegalStateException.class, () -> record.headers().add(new RecordHeader("test", "test".getBytes())));
+        RecordHeaders recordHeaders = (RecordHeaders) record.headers();
+        assertTrue(recordHeaders.isReadOnly());
+        assertThrows(IllegalStateException.class, () -> recordHeaders.add(new RecordHeader("test", "test".getBytes())));
 
         //ensure existing headers are not changed, and last header for key is still original value
         assertArrayEquals(record.headers().lastHeader("test").value(), "header2".getBytes());
@@ -1130,6 +1135,28 @@ private <T extends Serializer<String>> void doTestHeaders(Class<T> serializerCla
         producer.close(Duration.ofMillis(0));
     }
 
+    @Test
+    public void testHeadersFailure() {
+        Properties props = new Properties();
+        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, 5);
+        props.put(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, ProducerInterceptorForHeaders.class.getName());
+        Serializer<String> keySerializer = mock(StringSerializer.class);
+        Serializer<String> valueSerializer = mock(StringSerializer.class);
+
+        KafkaProducer<String, String> producer = new KafkaProducer<>(props, keySerializer, valueSerializer);
+        ProducerRecord<String, String> record = new ProducerRecord<>("topic", "key", "value");
+        Future<RecordMetadata> future = producer.send(record, (recordMetadata, exception) -> { });
+        try {
+            TestUtils.assertFutureThrows(TimeoutException.class, future);
+            //ensure headers are writable if send failure
+            RecordHeaders recordHeaders = (RecordHeaders) record.headers();
+            assertFalse(recordHeaders.isReadOnly());
+        } finally {
+            producer.close(Duration.ofMillis(0));
+        }
+    }
+
     @Test
     public void closeShouldBeIdempotent() {
         Properties producerProps = new Properties();
@@ -2500,6 +2527,29 @@ public void configure(Map<String, ?> configs) {
         }
     }
 
+    public static class ProducerInterceptorForHeaders implements ProducerInterceptor<byte[], byte[]> {
+
+        @Override
+        public ProducerRecord<byte[], byte[]> onSend(ProducerRecord<byte[], byte[]> record) {
+            return record;
+        }
+
+        @Override
+        public void onAcknowledgement(RecordMetadata metadata, Exception exception, Headers headers) {
+            RecordHeaders recordHeaders = (RecordHeaders) headers;
+            // Ensure that the headers are read-only, no matter send success or send failure
+            assertTrue(recordHeaders.isReadOnly());
+        }
+
+        @Override
+        public void close() {
+        }
+
+        @Override
+        public void configure(Map<String, ?> configs) {
+        }
+    }
+
     public static class ProducerInterceptorForClientId implements ProducerInterceptor<byte[], byte[]> {
 
         @Override