DeadLetterPublishingRecoverer Improvements

- make it easier to enhance headers
- also polish the `ThreadLocal` in `KafkaUtils`

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/listener/DeadLetterPublishingRecoverer.java

@@ -53,8 +53,7 @@
  */
 public class DeadLetterPublishingRecoverer implements ConsumerRecordRecoverer {
 
-	private static final LogAccessor LOGGER =
-			new LogAccessor(LogFactory.getLog(DeadLetterPublishingRecoverer.class));
+	protected final LogAccessor logger = new LogAccessor(LogFactory.getLog(getClass())); // NOSONAR
 
 	private static final BiFunction<ConsumerRecord<?, ?>, Exception, TopicPartition>
 		DEFAULT_DESTINATION_RESOLVER = (cr, e) -> new TopicPartition(cr.topic() + ".DLT", cr.partition());
@@ -69,6 +68,8 @@ public class DeadLetterPublishingRecoverer implements ConsumerRecordRecoverer {
 
 	private boolean retainExceptionHeader;
 
+	private BiFunction<ConsumerRecord<?, ?>, Exception, Headers> headersFunction = (rec, ex) -> null;
+
 	/**
 	 * Create an instance with the provided template and a default destination resolving
 	 * function that returns a TopicPartition based on the original topic (appended with ".DLT")
@@ -182,15 +183,26 @@ public void setRetainExceptionHeader(boolean retainExceptionHeader) {
 		this.retainExceptionHeader = retainExceptionHeader;
 	}
 
+	/**
+	 * Set a function which will be called to obtain additional headers to add to the
+	 * published record.
+	 * @param headersFunction the headers function.
+	 * @since 2.5.4
+	 */
+	public void setHeadersFunction(BiFunction<ConsumerRecord<?, ?>, Exception, Headers> headersFunction) {
+		Assert.notNull(headersFunction, "'headersFunction' cannot be null");
+		this.headersFunction = headersFunction;
+	}
+
 	@Override
 	public void accept(ConsumerRecord<?, ?> record, Exception exception) {
 		TopicPartition tp = this.destinationResolver.apply(record, exception);
 		boolean isKey = false;
 		DeserializationException deserEx = ListenerUtils.getExceptionFromHeader(record,
-				ErrorHandlingDeserializer.VALUE_DESERIALIZER_EXCEPTION_HEADER, LOGGER);
+				ErrorHandlingDeserializer.VALUE_DESERIALIZER_EXCEPTION_HEADER, this.logger);
 		if (deserEx == null) {
 			deserEx = ListenerUtils.getExceptionFromHeader(record,
-					ErrorHandlingDeserializer.KEY_DESERIALIZER_EXCEPTION_HEADER, LOGGER);
+					ErrorHandlingDeserializer.KEY_DESERIALIZER_EXCEPTION_HEADER, this.logger);
 			isKey = true;
 		}
 		Headers headers;
@@ -236,7 +248,7 @@ private KafkaOperations<Object, Object> findTemplateForValue(@Nullable Object va
 		if (key.isPresent()) {
 			return (KafkaOperations<Object, Object>) this.templates.get(key.get());
 		}
-		LOGGER.warn(() -> "Failed to find a template for " + value.getClass() + " attempting to use the last entry");
+		this.logger.warn(() -> "Failed to find a template for " + value.getClass() + " attempting to use the last entry");
 		return (KafkaOperations<Object, Object>) this.templates.values()
 				.stream()
 				.reduce((first,  second) -> second)
@@ -276,13 +288,13 @@ protected ProducerRecord<Object, Object> createProducerRecord(ConsumerRecord<?,
 	protected void publish(ProducerRecord<Object, Object> outRecord, KafkaOperations<Object, Object> kafkaTemplate) {
 		try {
 			kafkaTemplate.send(outRecord).addCallback(result -> {
-				LOGGER.debug(() -> "Successful dead-letter publication: " + result);
+				this.logger.debug(() -> "Successful dead-letter publication: " + result);
 			}, ex -> {
-				LOGGER.error(ex, () -> "Dead-letter publication failed for: " + outRecord);
+				this.logger.error(ex, () -> "Dead-letter publication failed for: " + outRecord);
 			});
 		}
 		catch (Exception e) {
-			LOGGER.error(e, () -> "Dead-letter publication failed for: " + outRecord);
+			this.logger.error(e, () -> "Dead-letter publication failed for: " + outRecord);
 		}
 	}
 
@@ -306,6 +318,10 @@ private void enhanceHeaders(Headers kafkaHeaders, ConsumerRecord<?, ?> record, E
 		}
 		kafkaHeaders.add(new RecordHeader(KafkaHeaders.DLT_EXCEPTION_STACKTRACE,
 				getStackTraceAsString(exception).getBytes(StandardCharsets.UTF_8)));
+		Headers headers = this.headersFunction.apply(record, exception);
+		if (headers != null) {
+			headers.forEach(header -> kafkaHeaders.add(header));
+		}
 	}
 
 	private String getStackTraceAsString(Throwable cause) {

spring-kafka/src/main/java/org/springframework/kafka/support/KafkaUtils.java

@@ -40,7 +40,7 @@ public final class KafkaUtils {
 	public static final boolean MICROMETER_PRESENT = ClassUtils.isPresent(
 			"io.micrometer.core.instrument.MeterRegistry", KafkaUtils.class.getClassLoader());
 
-	private static ThreadLocal<String> groupIds = new ThreadLocal<>();
+	private static final ThreadLocal<String> GROUP_IDS = new ThreadLocal<>();
 
 	/**
 	 * Return true if the method return type is {@link Message} or
@@ -78,7 +78,7 @@ public static boolean returnTypeMessageOrCollectionOf(Method method) {
 	 * @since 2.3
 	 */
 	public static void setConsumerGroupId(String groupId) {
-		KafkaUtils.groupIds.set(groupId);
+		KafkaUtils.GROUP_IDS.set(groupId);
 	}
 
 	/**
@@ -87,15 +87,15 @@ public static void setConsumerGroupId(String groupId) {
 	 * @since 2.3
 	 */
 	public static String getConsumerGroupId() {
-		return KafkaUtils.groupIds.get();
+		return KafkaUtils.GROUP_IDS.get();
 	}
 
 	/**
 	 * Clear the group id for the consumer bound to this thread.
 	 * @since 2.3
 	 */
 	public static void clearConsumerGroupId() {
-		KafkaUtils.groupIds.remove();
+		KafkaUtils.GROUP_IDS.remove();
 	}
 
 	private KafkaUtils() {

spring-kafka/src/test/java/org/springframework/kafka/listener/DeadLetterPublishingRecovererTests.java

@@ -124,6 +124,9 @@ void valueHeaderStripped() {
 		Headers headers = new RecordHeaders();
 		headers.add(new RecordHeader(ErrorHandlingDeserializer.VALUE_DESERIALIZER_EXCEPTION_HEADER, header(false)));
 		headers.add(new RecordHeader(ErrorHandlingDeserializer.KEY_DESERIALIZER_EXCEPTION_HEADER, header(true)));
+		Headers custom = new RecordHeaders();
+		custom.add(new RecordHeader("foo", "bar".getBytes()));
+		recoverer.setHeadersFunction((rec, ex) -> custom);
 		willReturn(new SettableListenableFuture<Object>()).given(template).send(any(ProducerRecord.class));
 		ConsumerRecord<String, String> record = new ConsumerRecord<>("foo", 0, 0L, 0L, TimestampType.CREATE_TIME,
 				0L, 0, 0, "bar", "baz", headers);
@@ -133,6 +136,7 @@ void valueHeaderStripped() {
 		headers = captor.getValue().headers();
 		assertThat(headers.lastHeader(ErrorHandlingDeserializer.VALUE_DESERIALIZER_EXCEPTION_HEADER)).isNull();
 		assertThat(headers.lastHeader(ErrorHandlingDeserializer.KEY_DESERIALIZER_EXCEPTION_HEADER)).isNotNull();
+		assertThat(headers.lastHeader("foo")).isNotNull();
 	}
 
 	@SuppressWarnings({ "rawtypes", "unchecked" })

src/reference/asciidoc/kafka.adoc

@@ -4868,6 +4868,12 @@ The record sent to the dead-letter topic is enhanced with the following headers:
 * `KafkaHeaders.DLT_ORIGINAL_TIMESTAMP`: The original timestamp.
 * `KafkaHeaders.DLT_ORIGINAL_TIMESTAMP_TYPE`: The original timestamp type.
 
+There are two mechanisms to add more headers.
+
+1. Subclass the recoverer and override `createProducerRecord()` - call `super.createProducerRecord()` and add more headers.
+2. Provide a `BiFunction` to receive the consumer record and exception, returning a `Headers` object; headers from there will be copied to the final producer record.
+
+The second is simpler to implement but the first has more information, including the already assembled standard headers.
 
 Starting with version 2.3, when used in conjunction with an `ErrorHandlingDeserializer`, the publisher will restore the record `value()`, in the dead-letter producer record, to the original value that failed to be deserialized.
 Previously, the `value()` was null and user code had to decode the `DeserializationException` from the message headers.
@@ -4903,7 +4909,7 @@ By default, these headers are not retained in the message published to the dead
 In that case, the `DLT_*` headers are based on the value deserialization and the key `DeserializationException` is retained in the header.
 
 [[kerberos]]
-==== Kerberos
+==== JAAS and Kerberos
 
 Starting with version 2.0, a `KafkaJaasLoginModuleInitializer` class has been added to assist with Kerberos configuration.
 You can add this bean, with the desired configuration, to your application context.