GH-763: Add DlqPartitionFunction

Resolves https://github.com/spring-cloud/spring-cloud-stream-binder-kafka/issues/763

Allow users to select the DLQ partition based on

* consumer group
* consumer record (which includes the original topic/partition)
* exception

Kafka Streams DLQ docs changes

docs/src/main/asciidoc/dlq.adoc

@@ -1,7 +1,32 @@
 [[kafka-dlq-processing]]
 === Dead-Letter Topic Processing
 
-Because you cannot anticipate how users would want to dispose of dead-lettered messages, the framework does not provide any standard mechanism to handle them.
+[[dlq-partition-selection]]
+==== Dead-Letter Topic Partition Selection
+
+By default, records are published to the Dead-Letter topic using the same partition as the original record.
+This means the Dead-Letter topic must have at least as many partitions as the original record.
+
+To change this behavior, add a `DlqPartitionFunction` implementation as a `@Bean` to the application context.
+Only one such bean can be present.
+The function is provided with the consumer group, the failed `ConsumerRecord` and the exception.
+For example, if you always with to route to partition 0, you might use:
+
+====
+[source, java]
+----
+@Bean
+public DlqPartitionFunction partitionFunction() {
+    return (group, record, ex) -> 0;
+}
+----
+====
+
+
+[[dlq-handling]]
+==== Handling Records in a Dead-Letter Topic
+
+Because the framework cannot anticipate how users would want to dispose of dead-lettered messages, it does not provide any standard mechanism to handle them.
 If the reason for the dead-lettering is transient, you may wish to route the messages back to the original topic.
 However, if the problem is a permanent issue, that could cause an infinite loop.
 The sample Spring Boot application within this topic is an example of how to route those messages back to the original topic, but it moves them to a "`parking lot`" topic after three attempts.

docs/src/main/asciidoc/kafka-streams.adoc

@@ -882,6 +882,23 @@ spring.cloud.stream.kafka.streams.bindings.input.consumer.dlqName: custom-dlq
 If this is set, then the error records are sent to the topic `custom-dlq`. If this is not set, then it will create a DLQ
 topic with the name `error.<input-topic-name>.<group-name>`.
 
+By default, records are published to the Dead-Letter topic using the same partition as the original record.
+This means the Dead-Letter topic must have at least as many partitions as the original record.
+
+To change this behavior, add a `DlqPartitionFunction` implementation as a `@Bean` to the application context.
+Only one such bean can be present.
+The function is provided with the consumer group, the failed `ConsumerRecord` and the exception.
+For example, if you always with to route to partition 0, you might use:
+
+
+[source, java]
+----
+@Bean
+public DlqPartitionFunction partitionFunction() {
+    return (group, record, ex) -> 0;
+}
+----
+
 A couple of things to keep in mind when using the exception handling feature in Kafka Streams binder.
 
 * The property `spring.cloud.stream.kafka.streams.binder.serdeError` is applicable for the entire application. This implies

docs/src/main/asciidoc/overview.adoc

@@ -208,6 +208,8 @@ The DLQ topic name can be configurable by setting the `dlqName` property.
 This provides an alternative option to the more common Kafka replay scenario for the case when the number of errors is relatively small and replaying the entire original topic may be too cumbersome.
 See <<kafka-dlq-processing>> processing for more information.
 Starting with version 2.0, messages sent to the DLQ topic are enhanced with the following headers: `x-original-topic`, `x-exception-message`, and `x-exception-stacktrace` as `byte[]`.
+By default, a failed record is sent to the same partition number in the DLQ topic as the original record.
+See <<dlq-partition-selection>> for how to change that behavior.
 **Not allowed when `destinationIsPattern` is `true`.**
 +
 Default: `false`.

spring-cloud-stream-binder-kafka-core/src/main/java/org/springframework/cloud/stream/binder/kafka/utils/DlqPartitionFunction.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2019-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.stream.binder.kafka.utils;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+import org.springframework.lang.Nullable;
+
+/**
+ * A TriFunction that takes a consumer group, consumer record, and throwable and returns
+ * which partition to publish to the dead letter topic. Returning {@code null} means Kafka
+ * will choose the partition.
+ *
+ * @author Gary Russell
+ * @since 3.0
+ *
+ */
+@FunctionalInterface
+public interface DlqPartitionFunction {
+
+	/**
+	 * Apply the function.
+	 * @param group the consumer group.
+	 * @param record the consumer record.
+	 * @param throwable the exception.
+	 * @return the DLQ partition, or null.
+	 */
+	@Nullable
+	Integer apply(String group, ConsumerRecord<?, ?> record, Throwable throwable);
+
+}

spring-cloud-stream-binder-kafka-streams/src/main/java/org/springframework/cloud/stream/binder/kafka/streams/KafkaStreamsBinderUtils.java

@@ -18,7 +18,9 @@ package org.springframework.cloud.stream.binder.kafka.streams;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.BiFunction;
 
+import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
@@ -32,6 +34,7 @@ import org.springframework.cloud.stream.binder.kafka.properties.KafkaProducerPro
 import org.springframework.cloud.stream.binder.kafka.provisioning.KafkaTopicProvisioner;
 import org.springframework.cloud.stream.binder.kafka.streams.properties.KafkaStreamsBinderConfigurationProperties;
 import org.springframework.cloud.stream.binder.kafka.streams.properties.KafkaStreamsConsumerProperties;
+import org.springframework.cloud.stream.binder.kafka.utils.DlqPartitionFunction;
 import org.springframework.context.ApplicationContext;
 import org.springframework.core.MethodParameter;
 import org.springframework.kafka.core.DefaultKafkaProducerFactory;
@@ -45,6 +48,7 @@ import org.springframework.util.StringUtils;
  * Common methods used by various Kafka Streams types across the binders.
  *
  * @author Soby Chacko
+ * @author Gary Russell
  */
 final class KafkaStreamsBinderUtils {
 
@@ -56,6 +60,7 @@ final class KafkaStreamsBinderUtils {
 			ApplicationContext context, KafkaTopicProvisioner kafkaTopicProvisioner,
 			KafkaStreamsBinderConfigurationProperties binderConfigurationProperties,
 			ExtendedConsumerProperties<KafkaStreamsConsumerProperties> properties) {
+
 		ExtendedConsumerProperties<KafkaConsumerProperties> extendedConsumerProperties = new ExtendedConsumerProperties<>(
 				properties.getExtension());
 		if (binderConfigurationProperties
@@ -71,6 +76,12 @@ final class KafkaStreamsBinderUtils {
 
 		if (extendedConsumerProperties.getExtension().isEnableDlq()) {
 
+			Map<String, DlqPartitionFunction> partitionFunctions =
+					context.getBeansOfType(DlqPartitionFunction.class, false, false);
+			DlqPartitionFunction partitionFunction = partitionFunctions.size() == 1
+					? partitionFunctions.values().iterator().next()
+					: (grp, rec, ex) -> rec.partition();
+
 			ProducerFactory<byte[], byte[]> producerFactory = getProducerFactory(
 					new ExtendedProducerProperties<>(
 							extendedConsumerProperties.getExtension().getDlqProducerProperties()),
@@ -78,15 +89,20 @@ final class KafkaStreamsBinderUtils {
 			KafkaTemplate<byte[], byte[]> kafkaTemplate = new KafkaTemplate<>(producerFactory);
 
 
+			BiFunction<ConsumerRecord<?, ?>, Exception, TopicPartition> destinationResolver =
+					(cr, e) -> new TopicPartition(extendedConsumerProperties.getExtension().getDlqName(),
+							partitionFunction.apply(group, cr, e));
 			DeadLetterPublishingRecoverer kafkaStreamsBinderDlqRecoverer = !StringUtils
 					.isEmpty(extendedConsumerProperties.getExtension().getDlqName())
-					? new DeadLetterPublishingRecoverer(kafkaTemplate, (cr, e) -> new TopicPartition(extendedConsumerProperties.getExtension()
-					.getDlqName(), cr.partition()))
+					? new DeadLetterPublishingRecoverer(kafkaTemplate, destinationResolver)
 					: null;
 			for (String inputTopic : inputTopics) {
 				if (StringUtils.isEmpty(
 						extendedConsumerProperties.getExtension().getDlqName())) {
-					kafkaStreamsBinderDlqRecoverer = new DeadLetterPublishingRecoverer(kafkaTemplate, (cr, e) -> new TopicPartition("error." + inputTopic + "." + group, cr.partition()));
+					destinationResolver = (cr, e) -> new TopicPartition("error." + inputTopic + "." + group,
+									partitionFunction.apply(group, cr, e));
+					kafkaStreamsBinderDlqRecoverer = new DeadLetterPublishingRecoverer(kafkaTemplate,
+							destinationResolver);
 				}
 
 				SendToDlqAndContinue sendToDlqAndContinue = context

spring-cloud-stream-binder-kafka-streams/src/test/java/org/springframework/cloud/stream/binder/kafka/streams/integration/DeserializationErrorHandlerByKafkaTests.java

@@ -40,6 +40,8 @@ import org.springframework.boot.test.mock.mockito.SpyBean;
 import org.springframework.cloud.stream.annotation.EnableBinding;
 import org.springframework.cloud.stream.annotation.StreamListener;
 import org.springframework.cloud.stream.binder.kafka.streams.annotations.KafkaStreamsProcessor;
+import org.springframework.cloud.stream.binder.kafka.utils.DlqPartitionFunction;
+import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.PropertySource;
 import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
 import org.springframework.kafka.core.DefaultKafkaProducerFactory;
@@ -67,7 +69,11 @@ public abstract class DeserializationErrorHandlerByKafkaTests {
 
 	@ClassRule
 	public static EmbeddedKafkaRule embeddedKafkaRule = new EmbeddedKafkaRule(1, true,
-			"DeserializationErrorHandlerByKafkaTests-out", "error.DeserializationErrorHandlerByKafkaTests-In.group", "error.word1.groupx", "error.word2.groupx");
+			"DeserializationErrorHandlerByKafkaTests-In",
+			"DeserializationErrorHandlerByKafkaTests-out",
+			"error.DeserializationErrorHandlerByKafkaTests-In.group",
+			"error.word1.groupx",
+			"error.word2.groupx");
 
 	private static EmbeddedKafkaBroker embeddedKafka = embeddedKafkaRule
 			.getEmbeddedKafka();
@@ -78,7 +84,7 @@ public abstract class DeserializationErrorHandlerByKafkaTests {
 	private static Consumer<String, String> consumer;
 
 	@BeforeClass
-	public static void setUp() throws Exception {
+	public static void setUp() {
 		System.setProperty("spring.cloud.stream.kafka.streams.binder.brokers",
 				embeddedKafka.getBrokersAsString());
 
@@ -112,14 +118,13 @@ public abstract class DeserializationErrorHandlerByKafkaTests {
 			extends DeserializationErrorHandlerByKafkaTests {
 
 		@Test
-		@SuppressWarnings("unchecked")
-		public void test() throws Exception {
+		public void test() {
 			Map<String, Object> senderProps = KafkaTestUtils.producerProps(embeddedKafka);
 			DefaultKafkaProducerFactory<Integer, String> pf = new DefaultKafkaProducerFactory<>(
 					senderProps);
 			KafkaTemplate<Integer, String> template = new KafkaTemplate<>(pf, true);
 			template.setDefaultTopic("DeserializationErrorHandlerByKafkaTests-In");
-			template.sendDefault("foobar");
+			template.sendDefault(1, null, "foobar");
 
 			Map<String, Object> consumerProps = KafkaTestUtils.consumerProps("foobar",
 					"false", embeddedKafka);
@@ -131,7 +136,8 @@ public abstract class DeserializationErrorHandlerByKafkaTests {
 
 			ConsumerRecord<String, String> cr = KafkaTestUtils.getSingleRecord(consumer1,
 					"error.DeserializationErrorHandlerByKafkaTests-In.group");
-			assertThat(cr.value().equals("foobar")).isTrue();
+			assertThat(cr.value()).isEqualTo("foobar");
+			assertThat(cr.partition()).isEqualTo(0); // custom partition function
 
 			// Ensuring that the deserialization was indeed done by Kafka natively
 			verify(conversionDelegate, never()).deserializeOnInbound(any(Class.class),
@@ -153,8 +159,7 @@ public abstract class DeserializationErrorHandlerByKafkaTests {
 			extends DeserializationErrorHandlerByKafkaTests {
 
 		@Test
-		@SuppressWarnings("unchecked")
-		public void test() throws Exception {
+		public void test() {
 			Map<String, Object> senderProps = KafkaTestUtils.producerProps(embeddedKafka);
 			DefaultKafkaProducerFactory<Integer, String> pf = new DefaultKafkaProducerFactory<>(
 					senderProps);
@@ -176,10 +181,10 @@ public abstract class DeserializationErrorHandlerByKafkaTests {
 
 			ConsumerRecord<String, String> cr1 = KafkaTestUtils.getSingleRecord(consumer1,
 					"error.word1.groupx");
-			assertThat(cr1.value().equals("foobar")).isTrue();
+			assertThat(cr1.value()).isEqualTo("foobar");
 			ConsumerRecord<String, String> cr2 = KafkaTestUtils.getSingleRecord(consumer1,
 					"error.word2.groupx");
-			assertThat(cr2.value().equals("foobar")).isTrue();
+			assertThat(cr2.value()).isEqualTo("foobar");
 
 			// Ensuring that the deserialization was indeed done by Kafka natively
 			verify(conversionDelegate, never()).deserializeOnInbound(any(Class.class),
@@ -208,6 +213,11 @@ public abstract class DeserializationErrorHandlerByKafkaTests {
 							"Count for " + key.key() + " : " + value));
 		}
 
+		@Bean
+		public DlqPartitionFunction partitionFunction() {
+			return (group, rec, ex) -> 0;
+		}
+
 	}
 
 }

spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaMessageChannelBinder.java

@@ -70,6 +70,7 @@ import org.springframework.cloud.stream.binder.kafka.properties.KafkaConsumerPro
 import org.springframework.cloud.stream.binder.kafka.properties.KafkaExtendedBindingProperties;
 import org.springframework.cloud.stream.binder.kafka.properties.KafkaProducerProperties;
 import org.springframework.cloud.stream.binder.kafka.provisioning.KafkaTopicProvisioner;
+import org.springframework.cloud.stream.binder.kafka.utils.DlqPartitionFunction;
 import org.springframework.cloud.stream.binding.MessageConverterConfigurer.PartitioningInterceptor;
 import org.springframework.cloud.stream.config.ListenerContainerCustomizer;
 import org.springframework.cloud.stream.config.MessageSourceCustomizer;
@@ -199,6 +200,8 @@ public class KafkaMessageChannelBinder extends
 
 	private final KafkaBindingRebalanceListener rebalanceListener;
 
+	private final DlqPartitionFunction dlqPartitionFunction;
+
 	private ProducerListener<byte[], byte[]> producerListener;
 
 	private KafkaExtendedBindingProperties extendedBindingProperties = new KafkaExtendedBindingProperties();
@@ -207,7 +210,7 @@ public class KafkaMessageChannelBinder extends
 			KafkaBinderConfigurationProperties configurationProperties,
 			KafkaTopicProvisioner provisioningProvider) {
 
-		this(configurationProperties, provisioningProvider, null, null, null);
+		this(configurationProperties, provisioningProvider, null, null, null, null);
 	}
 
 	public KafkaMessageChannelBinder(
@@ -216,7 +219,7 @@ public class KafkaMessageChannelBinder extends
 			ListenerContainerCustomizer<AbstractMessageListenerContainer<?, ?>> containerCustomizer,
 			KafkaBindingRebalanceListener rebalanceListener) {
 
-		this(configurationProperties, provisioningProvider, containerCustomizer, null, rebalanceListener);
+		this(configurationProperties, provisioningProvider, containerCustomizer, null, rebalanceListener, null);
 	}
 
 	public KafkaMessageChannelBinder(
@@ -224,7 +227,8 @@ public class KafkaMessageChannelBinder extends
 			KafkaTopicProvisioner provisioningProvider,
 			ListenerContainerCustomizer<AbstractMessageListenerContainer<?, ?>> containerCustomizer,
 			MessageSourceCustomizer<KafkaMessageSource<?, ?>> sourceCustomizer,
-			KafkaBindingRebalanceListener rebalanceListener) {
+			KafkaBindingRebalanceListener rebalanceListener,
+			DlqPartitionFunction dlqPartitionFunction) {
 
 		super(headersToMap(configurationProperties), provisioningProvider,
 				containerCustomizer, sourceCustomizer);
@@ -240,6 +244,9 @@ public class KafkaMessageChannelBinder extends
 			this.transactionManager = null;
 		}
 		this.rebalanceListener = rebalanceListener;
+		this.dlqPartitionFunction = dlqPartitionFunction != null
+				? dlqPartitionFunction
+				: (group, rec, ex) -> rec.partition();
 	}
 
 	private static String[] headersToMap(
@@ -1008,9 +1015,10 @@ public class KafkaMessageChannelBinder extends
 				Headers kafkaHeaders = new RecordHeaders(record.headers().toArray());
 				AtomicReference<ConsumerRecord<?, ?>> recordToSend = new AtomicReference<>(
 						record);
+				Throwable throwable = null;
 				if (message.getPayload() instanceof Throwable) {
 
-					Throwable throwable = (Throwable) message.getPayload();
+					throwable = (Throwable) message.getPayload();
 
 					HeaderMode headerMode = properties.getHeaderMode();
 
@@ -1074,7 +1082,8 @@ public class KafkaMessageChannelBinder extends
 				String dlqName = StringUtils.hasText(kafkaConsumerProperties.getDlqName())
 						? kafkaConsumerProperties.getDlqName()
 						: "error." + record.topic() + "." + group;
-				dlqSender.sendToDlq(recordToSend.get(), kafkaHeaders, dlqName);
+				dlqSender.sendToDlq(recordToSend.get(), kafkaHeaders, dlqName, group, throwable,
+						this.dlqPartitionFunction);
 			};
 		}
 		return null;
@@ -1327,11 +1336,12 @@ public class KafkaMessageChannelBinder extends
 
 		@SuppressWarnings("unchecked")
 		void sendToDlq(ConsumerRecord<?, ?> consumerRecord, Headers headers,
-				String dlqName) {
+				String dlqName, String group, Throwable throwable, DlqPartitionFunction partitionFunction) {
 			K key = (K) consumerRecord.key();
 			V value = (V) consumerRecord.value();
 			ProducerRecord<K, V> producerRecord = new ProducerRecord<>(dlqName,
-					consumerRecord.partition(), key, value, headers);
+					partitionFunction.apply(group, consumerRecord, throwable),
+					key, value, headers);
 
 			StringBuilder sb = new StringBuilder().append(" a message with key='")
 					.append(toDisplayString(ObjectUtils.nullSafeToString(key), 50))

spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/config/KafkaBinderConfiguration.java

@@ -39,6 +39,7 @@ import org.springframework.cloud.stream.binder.kafka.properties.JaasLoginModuleC
 import org.springframework.cloud.stream.binder.kafka.properties.KafkaBinderConfigurationProperties;
 import org.springframework.cloud.stream.binder.kafka.properties.KafkaExtendedBindingProperties;
 import org.springframework.cloud.stream.binder.kafka.provisioning.KafkaTopicProvisioner;
+import org.springframework.cloud.stream.binder.kafka.utils.DlqPartitionFunction;
 import org.springframework.cloud.stream.config.ListenerContainerCustomizer;
 import org.springframework.cloud.stream.config.MessageSourceCustomizer;
 import org.springframework.context.ApplicationContext;
@@ -103,11 +104,13 @@ public class KafkaBinderConfiguration {
 			KafkaTopicProvisioner provisioningProvider,
 			@Nullable ListenerContainerCustomizer<AbstractMessageListenerContainer<?, ?>> listenerContainerCustomizer,
 			@Nullable MessageSourceCustomizer<KafkaMessageSource<?, ?>> sourceCustomizer,
-			ObjectProvider<KafkaBindingRebalanceListener> rebalanceListener) {
+			ObjectProvider<KafkaBindingRebalanceListener> rebalanceListener,
+			ObjectProvider<DlqPartitionFunction> dlqPartitionFunction) {
 
 		KafkaMessageChannelBinder kafkaMessageChannelBinder = new KafkaMessageChannelBinder(
 				configurationProperties, provisioningProvider,
-				listenerContainerCustomizer, sourceCustomizer, rebalanceListener.getIfUnique());
+				listenerContainerCustomizer, sourceCustomizer, rebalanceListener.getIfUnique(),
+				dlqPartitionFunction.getIfUnique());
 		kafkaMessageChannelBinder.setProducerListener(this.producerListener);
 		kafkaMessageChannelBinder
 				.setExtendedBindingProperties(this.kafkaExtendedBindingProperties);

spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaBinderTests.java

@@ -88,6 +88,7 @@ import org.springframework.cloud.stream.binder.kafka.properties.KafkaBinderConfi
 import org.springframework.cloud.stream.binder.kafka.properties.KafkaConsumerProperties;
 import org.springframework.cloud.stream.binder.kafka.properties.KafkaProducerProperties;
 import org.springframework.cloud.stream.binder.kafka.provisioning.KafkaTopicProvisioner;
+import org.springframework.cloud.stream.binder.kafka.utils.DlqPartitionFunction;
 import org.springframework.cloud.stream.binder.kafka.utils.KafkaTopicUtils;
 import org.springframework.cloud.stream.binding.MessageConverterConfigurer.PartitioningInterceptor;
 import org.springframework.cloud.stream.config.BindingProperties;
@@ -211,8 +212,16 @@ public class KafkaBinderTests extends
 		return binder;
 	}
 
-	private Binder getBinder(
+	private KafkaTestBinder getBinder(
 			KafkaBinderConfigurationProperties kafkaBinderConfigurationProperties) {
+
+		return getBinder(kafkaBinderConfigurationProperties, null);
+	}
+
+	private KafkaTestBinder getBinder(
+			KafkaBinderConfigurationProperties kafkaBinderConfigurationProperties,
+			DlqPartitionFunction dlqPartitionFunction) {
+
 		KafkaTopicProvisioner provisioningProvider = new KafkaTopicProvisioner(
 				kafkaBinderConfigurationProperties, new TestKafkaProperties());
 		try {
@@ -222,7 +231,7 @@ public class KafkaBinderTests extends
 			throw new RuntimeException(e);
 		}
 		return new KafkaTestBinder(kafkaBinderConfigurationProperties,
-				provisioningProvider);
+				provisioningProvider, dlqPartitionFunction);
 	}
 
 	private KafkaBinderConfigurationProperties createConfigurationProperties() {
@@ -869,7 +878,9 @@ public class KafkaBinderTests extends
 	}
 
 	private void testDlqGuts(boolean withRetry, HeaderMode headerMode) throws Exception {
-		AbstractKafkaTestBinder binder = getBinder();
+		KafkaBinderConfigurationProperties binderConfig = createConfigurationProperties();
+		binderConfig.setMinPartitionCount(2);
+		AbstractKafkaTestBinder binder = getBinder(binderConfig, (group, rec, ex) -> 0);
 
 		ExtendedProducerProperties<KafkaProducerProperties> producerProperties = createProducerProperties();
 		producerProperties.getExtension()
@@ -907,7 +918,7 @@ public class KafkaBinderTests extends
 		MessageListenerContainer container = TestUtils.getPropertyValue(consumerBinding,
 				"lifecycle.messageListenerContainer", MessageListenerContainer.class);
 		assertThat(container.getContainerProperties().getTopicPartitions().length)
-				.isEqualTo(2);
+				.isEqualTo(4); // 2 topics 2 partitions each
 
 		ExtendedConsumerProperties<KafkaConsumerProperties> dlqConsumerProperties = createConsumerProperties();
 		dlqConsumerProperties.setMaxAttempts(1);
@@ -936,7 +947,9 @@ public class KafkaBinderTests extends
 		binderBindUnbindLatency();
 		String testMessagePayload = "test." + UUID.randomUUID().toString();
 		Message<byte[]> testMessage = MessageBuilder
-				.withPayload(testMessagePayload.getBytes()).build();
+				.withPayload(testMessagePayload.getBytes())
+				.setHeader(KafkaHeaders.PARTITION_ID, 1)
+				.build();
 		moduleOutputChannel.send(testMessage);
 
 		Message<?> receivedMessage = receive(dlqChannel, 3);
@@ -951,7 +964,7 @@ public class KafkaBinderTests extends
 							.isEqualTo(producerName);
 
 			assertThat(receivedMessage.getHeaders()
-					.get(KafkaMessageChannelBinder.X_ORIGINAL_PARTITION)).isEqualTo(0);
+					.get(KafkaMessageChannelBinder.X_ORIGINAL_PARTITION)).isEqualTo(1);
 
 			assertThat(receivedMessage.getHeaders()
 					.get(KafkaMessageChannelBinder.X_ORIGINAL_OFFSET)).isEqualTo(0);
@@ -971,6 +984,8 @@ public class KafkaBinderTests extends
 					.get(KafkaMessageChannelBinder.X_EXCEPTION_STACKTRACE)).isNotNull();
 			assertThat(receivedMessage.getHeaders()
 					.get(KafkaMessageChannelBinder.X_EXCEPTION_FQCN)).isNotNull();
+			assertThat(receivedMessage.getHeaders()
+					.get(KafkaHeaders.RECEIVED_PARTITION_ID)).isEqualTo(0);
 		}
 		else if (!HeaderMode.none.equals(headerMode)) {
 			assertThat(handler.getInvocationCount())
@@ -982,7 +997,7 @@ public class KafkaBinderTests extends
 
 			assertThat(receivedMessage.getHeaders()
 					.get(KafkaMessageChannelBinder.X_ORIGINAL_PARTITION)).isEqualTo(
-							ByteBuffer.allocate(Integer.BYTES).putInt(0).array());
+							ByteBuffer.allocate(Integer.BYTES).putInt(1).array());
 
 			assertThat(receivedMessage.getHeaders()
 					.get(KafkaMessageChannelBinder.X_ORIGINAL_OFFSET)).isEqualTo(
@@ -1005,6 +1020,9 @@ public class KafkaBinderTests extends
 
 			assertThat(receivedMessage.getHeaders()
 					.get(KafkaMessageChannelBinder.X_EXCEPTION_FQCN)).isNotNull();
+
+			assertThat(receivedMessage.getHeaders()
+					.get(KafkaHeaders.RECEIVED_PARTITION_ID)).isEqualTo(0);
 		}
 		else {
 			assertThat(receivedMessage.getHeaders()

spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaTestBinder.java

@@ -20,6 +20,7 @@ import org.springframework.cloud.stream.binder.ExtendedConsumerProperties;
 import org.springframework.cloud.stream.binder.kafka.properties.KafkaBinderConfigurationProperties;
 import org.springframework.cloud.stream.binder.kafka.properties.KafkaConsumerProperties;
 import org.springframework.cloud.stream.binder.kafka.provisioning.KafkaTopicProvisioner;
+import org.springframework.cloud.stream.binder.kafka.utils.DlqPartitionFunction;
 import org.springframework.cloud.stream.provisioning.ConsumerDestination;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Configuration;
@@ -38,12 +39,19 @@ import org.springframework.kafka.support.ProducerListener;
  */
 public class KafkaTestBinder extends AbstractKafkaTestBinder {
 
-	@SuppressWarnings({ "rawtypes", "unchecked" })
 	KafkaTestBinder(KafkaBinderConfigurationProperties binderConfiguration,
 			KafkaTopicProvisioner kafkaTopicProvisioner) {
+
+		this(binderConfiguration, kafkaTopicProvisioner, null);
+	}
+
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	KafkaTestBinder(KafkaBinderConfigurationProperties binderConfiguration,
+			KafkaTopicProvisioner kafkaTopicProvisioner, DlqPartitionFunction dlqPartitionFunction) {
+
 		try {
 			KafkaMessageChannelBinder binder = new KafkaMessageChannelBinder(
-					binderConfiguration, kafkaTopicProvisioner) {
+					binderConfiguration, kafkaTopicProvisioner, null, null, null, dlqPartitionFunction) {
 
 				/*
 				 * Some tests use multiple instance indexes for the same topic; we need to