GH-2438: RetryTopic Destination Partition Select

Resolves #2438

**cherry-pick to 2.9.x**

Author: garyrussell

spring-kafka-docs/src/main/asciidoc/retrytopic.adoc

@@ -200,6 +200,33 @@ public class MyRetryTopicConfiguration extends RetryTopicConfigurationSupport {
 IMPORTANT: When using this configuration approach, the `@EnableKafkaRetryTopic` annotation should not be used to prevent context failing to start due to duplicated beans.
 Use the simple `@EnableKafka` annotation instead.
 
+When `autoCreateTopics` is true, the main and retry topics will be created with the specified number of partitions and replication factor.
+To override these values for a particular topic (e.g. the main topic or DLT), simply add a `NewTopic` `@Bean` with the required properties; that will override the auto creation properties.
+
+IMPORTANT: By default, records are published to the retry topic(s) using the original partition of the received record.
+If the retry topics have fewer partitions than the main topic, you should configure the framework appropriately; an example follows.
+
+====
+[source, java]
+----
+@EnableKafka
+@Configuration
+public class Config extends RetryTopicConfigurationSupport {
+
+    @Override
+    protected Consumer<DeadLetterPublishingRecovererFactory> configureDeadLetterPublishingContainerFactory() {
+        return dlprf -> dlprf.setPartitionResolver((cr, nextTopic) -> null);
+    }
+
+    ...
+
+}
+----
+====
+
+The parameters to the function are the consumer record and the name of the next topic.
+You can return a specific partition number, or `null` to indicate that the `KafkaProducer` should determine the partition.
+
 ==== Features
 
 Most of the features are available both for the `@RetryableTopic` annotation and the `RetryTopicConfiguration` beans.

spring-kafka/src/main/java/org/springframework/kafka/retrytopic/DeadLetterPublishingRecovererFactory.java

@@ -69,6 +69,8 @@ public class DeadLetterPublishingRecovererFactory {
 
 	private ListenerExceptionLoggingStrategy loggingStrategy = ListenerExceptionLoggingStrategy.AFTER_RETRIES_EXHAUSTED;
 
+	private BiFunction<ConsumerRecord<?, ?>, String, Integer> partitionResolver = (cr, nextTopic) -> cr.partition();
+
 	public DeadLetterPublishingRecovererFactory(DestinationTopicResolver destinationTopicResolver) {
 		this.destinationTopicResolver = destinationTopicResolver;
 	}
@@ -83,6 +85,19 @@ public void setHeadersFunction(BiFunction<ConsumerRecord<?, ?>, Exception, Heade
 		this.headersFunction = headersFunction;
 	}
 
+	/**
+	 * Set a resolver for the partition number to publish to. By default the same partition as
+	 * the consumer record is used. If the resolver returns {@code null} or a negative number, the
+	 * partition is set to null in the producer record and the {@code KafkaProducer} decides which
+	 * partition to publish to.
+	 * @param resolver the resolver.
+	 * @since 2.9.2
+	 */
+	public void setPartitionResolver(BiFunction<ConsumerRecord<?, ?>, String, Integer> resolver) {
+		Assert.notNull(resolver, "'resolver' cannot be null");
+		this.partitionResolver = resolver;
+	}
+
 	/**
 	 * Add exception type to the default list. By default, the following exceptions will
 	 * not be retried:
@@ -253,13 +268,17 @@ private static String getRecordInfo(ConsumerRecord<?, ?> cr) {
 	 * and if it doesn't it sets -1, to allow the Producer itself to assign a partition to the record.</p>
 	 *
 	 * <p>Subclasses can inherit from this method to override the implementation, if necessary.</p>
+	 * The destination partition can also be customized using {@link #setPartitionResolver(BiFunction)}.
 	 *
 	 * @param cr The original {@link ConsumerRecord}, which is to be forwarded to DLT
 	 * @param nextDestination The next {@link DestinationTopic}, where the consumerRecord is to be forwarded
-	 * @return An instance of {@link TopicPartition}, specifying the topic and partition, where the cr is to be sent
+	 * @return An instance of {@link TopicPartition}, specifying the topic and partition, where the cr is to be sent.
+	 * @see #setPartitionResolver(BiFunction)
 	 */
 	protected TopicPartition resolveTopicPartition(final ConsumerRecord<?, ?> cr, final DestinationTopic nextDestination) {
-		return new TopicPartition(nextDestination.getDestinationName(), cr.partition());
+		String nextTopic = nextDestination.getDestinationName();
+		Integer partition = this.partitionResolver.apply(cr, nextTopic);
+		return new TopicPartition(nextTopic, partition == null ? -1 : partition);
 	}
 
 	private int getAttempts(ConsumerRecord<?, ?> consumerRecord) {

spring-kafka/src/test/java/org/springframework/kafka/retrytopic/PartitionResolverTests.java

@@ -0,0 +1,137 @@
+/*
+ * Copyright 2022 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.kafka.retrytopic;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.willAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.kafka.annotation.EnableKafka;
+import org.springframework.kafka.annotation.KafkaListener;
+import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory;
+import org.springframework.kafka.core.ConsumerFactory;
+import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
+import org.springframework.kafka.core.DefaultKafkaProducerFactory;
+import org.springframework.kafka.core.KafkaOperations;
+import org.springframework.kafka.core.KafkaTemplate;
+import org.springframework.kafka.support.SendResult;
+import org.springframework.kafka.test.EmbeddedKafkaBroker;
+import org.springframework.kafka.test.context.EmbeddedKafka;
+import org.springframework.kafka.test.utils.KafkaTestUtils;
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
+
+/**
+ * @author Gary Russell
+ * @since 2.9.2
+ *
+ */
+@DirtiesContext
+@SpringJUnitConfig
+@EmbeddedKafka(topics = "partition.resolver.tests")
+public class PartitionResolverTests extends AbstractRetryTopicIntegrationTests {
+
+	@Test
+	void testNullPartition(@Autowired KafkaOperations<Integer, String> template,
+			@Autowired EmbeddedKafkaBroker broker, @Autowired Config config) throws InterruptedException {
+
+		Map<String, Object> producerProps = KafkaTestUtils.producerProps(broker);
+		DefaultKafkaProducerFactory<Integer, String> pf = new DefaultKafkaProducerFactory<>(producerProps);
+		KafkaTemplate<Integer, String> kt = new KafkaTemplate<>(pf);
+		kt.send("partition.resolver.tests", 1, null, "test");
+		assertThat(config.latch.await(10, TimeUnit.SECONDS)).isTrue();
+		@SuppressWarnings("unchecked")
+		ArgumentCaptor<ProducerRecord<Integer, String>> captor = ArgumentCaptor.forClass(ProducerRecord.class);
+		verify(template).send(captor.capture());
+		assertThat(captor.getValue().partition()).isNull();
+	}
+
+	@EnableKafka
+	@Configuration
+	public static class Config extends RetryTopicConfigurationSupport {
+
+		final CountDownLatch latch = new CountDownLatch(1);
+
+		@Override
+		protected Consumer<DeadLetterPublishingRecovererFactory> configureDeadLetterPublishingContainerFactory() {
+			return dlprf -> dlprf.setPartitionResolver((cr, nextTopic) -> null);
+		}
+
+		@Bean
+		RetryTopicConfiguration myRetryTopic(KafkaOperations<Integer, String> template) {
+			return RetryTopicConfigurationBuilder
+					.newInstance()
+					.create(template);
+		}
+
+		@SuppressWarnings("unchecked")
+		@Bean
+		KafkaOperations<Integer, String> template() {
+			KafkaOperations<Integer, String> mock = mock(KafkaOperations.class);
+			CompletableFuture<SendResult<Integer, String>> future = new CompletableFuture<>();
+			future.complete(mock(SendResult.class));
+			willAnswer(inv -> {
+				latch.countDown();
+				return future;
+			}).given(mock).send(any(ProducerRecord.class));
+			return mock;
+		}
+
+		@KafkaListener(topics = "partition.resolver.tests")
+		void listen(String in) {
+			throw new RuntimeException("test");
+		}
+
+		@Bean
+		ConsumerFactory<Integer, String> cf(EmbeddedKafkaBroker broker) {
+			Map<String, Object> props = KafkaTestUtils.consumerProps("prt", "false", broker);
+			return new DefaultKafkaConsumerFactory<>(props);
+		}
+
+		@Bean
+		ConcurrentKafkaListenerContainerFactory<Integer, String> kafkaListenerContainerFactory(
+				ConsumerFactory<Integer, String> cf) {
+			ConcurrentKafkaListenerContainerFactory<Integer, String> factory = new ConcurrentKafkaListenerContainerFactory<>();
+			factory.setConsumerFactory(cf);
+			return factory;
+		}
+
+		@Bean
+		TaskScheduler taskScheduler() {
+			return new ThreadPoolTaskScheduler();
+		}
+
+	}
+
+}