GH-1319: Agg. Replying Template improvement

Resolves #1319

- change the release strategy to a `BiConsumer`
- call the strategy during normal message delivery as before
- now call the strategy when releasing a partial result after timeout
- allow the strategy to modify the list (change from `Collection` to `List`).

**cherry-pick to 2.3.x**

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/requestreply/AggregatingReplyingKafkaTemplate.java

@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiPredicate;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -74,7 +75,7 @@ public class AggregatingReplyingKafkaTemplate<K, V, R>
 
 	private final Map<TopicPartition, Long> offsets = new HashMap<>();
 
-	private final Predicate<Collection<ConsumerRecord<K, R>>> releaseStrategy;
+	private final BiPredicate<List<ConsumerRecord<K, R>>, Boolean> releaseStrategy;
 
 	private Duration commitTimeout = Duration.ofSeconds(DEFAULT_COMMIT_TIMEOUT);
 
@@ -86,11 +87,32 @@ public class AggregatingReplyingKafkaTemplate<K, V, R>
 	 * @param producerFactory the producer factory.
 	 * @param replyContainer the reply container.
 	 * @param releaseStrategy the release strategy.
+	 * @deprecated in favor of
+	 * {@link #AggregatingReplyingKafkaTemplate(ProducerFactory, GenericMessageListenerContainer, BiPredicate)}
 	 */
+	@Deprecated
 	public AggregatingReplyingKafkaTemplate(ProducerFactory<K, V> producerFactory,
 			GenericMessageListenerContainer<K, Collection<ConsumerRecord<K, R>>> replyContainer,
 			Predicate<Collection<ConsumerRecord<K, R>>> releaseStrategy) {
 
+		this(producerFactory, replyContainer, (records, timeout) -> timeout || releaseStrategy.test(records));
+	}
+
+	/**
+	 * Construct an instance using the provided parameter arguments. The releaseStrategy
+	 * is consulted to determine when a collection is "complete".
+	 * @param producerFactory the producer factory.
+	 * @param replyContainer the reply container.
+	 * @param releaseStrategy the release strategy which is a {@link BiPredicate} which is
+	 * passed the current list and a boolean to indicate if this is for a normal delivery
+	 * or a timeout (when {@link #setReturnPartialOnTimeout(boolean)} is true. The
+	 * predicate may modify the list of records.
+	 * @since 2.3.5
+	 */
+	public AggregatingReplyingKafkaTemplate(ProducerFactory<K, V> producerFactory,
+			GenericMessageListenerContainer<K, Collection<ConsumerRecord<K, R>>> replyContainer,
+			BiPredicate<List<ConsumerRecord<K, R>>, Boolean> releaseStrategy) {
+
 		super(producerFactory, replyContainer);
 		Assert.notNull(releaseStrategy, "'releaseStrategy' cannot be null");
 		AckMode ackMode = replyContainer.getContainerProperties().getAckMode();
@@ -133,7 +155,7 @@ public void onMessage(List<ConsumerRecord<K, Collection<ConsumerRecord<K, R>>>>
 						List<ConsumerRecord<K, R>> list = addToCollection(record, correlationId).stream()
 								.map(RecordHolder::getRecord)
 								.collect(Collectors.toList());
-						if (this.releaseStrategy.test(list)) {
+						if (this.releaseStrategy.test(list, false)) {
 							ConsumerRecord<K, Collection<ConsumerRecord<K, R>>> done =
 									new ConsumerRecord<>(AGGREGATED_RESULTS_TOPIC, 0, 0L, null, list);
 							done.headers()
@@ -164,12 +186,12 @@ protected synchronized boolean handleTimeout(CorrelationKey correlationId,
 			List<ConsumerRecord<K, R>> list = removed.stream()
 					.map(RecordHolder::getRecord)
 					.collect(Collectors.toList());
-			future.set(new ConsumerRecord<>(PARTIAL_RESULTS_AFTER_TIMEOUT_TOPIC, 0, 0L, null, list));
-			return true;
-		}
-		else {
-			return false;
+			if (this.releaseStrategy.test(list, true)) {
+				future.set(new ConsumerRecord<>(PARTIAL_RESULTS_AFTER_TIMEOUT_TOPIC, 0, 0L, null, list));
+				return true;
+			}
 		}
+		return false;
 	}
 
 	private void checkOffsetsAndCommitIfNecessary(List<ConsumerRecord<K, R>> list, Consumer<?, ?> consumer) {

spring-kafka/src/test/java/org/springframework/kafka/requestreply/ReplyingKafkaTemplateTests.java

@@ -38,6 +38,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.kafka.clients.consumer.Consumer;
@@ -295,7 +296,7 @@ public void testGoodWithSimpleMapper() throws Exception {
 	@Test
 	public void testAggregateNormal() throws Exception {
 		AggregatingReplyingKafkaTemplate<Integer, String, String> template = aggregatingTemplate(
-				new TopicPartitionOffset(D_REPLY, 0), 2);
+				new TopicPartitionOffset(D_REPLY, 0), 2, new AtomicInteger());
 		try {
 			template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(D_REQUEST, null, null, null, "foo");
@@ -324,7 +325,7 @@ public void testAggregateNormal() throws Exception {
 	@Disabled("time sensitive")
 	public void testAggregateTimeout() throws Exception {
 		AggregatingReplyingKafkaTemplate<Integer, String, String> template = aggregatingTemplate(
-				new TopicPartitionOffset(E_REPLY, 0), 3);
+				new TopicPartitionOffset(E_REPLY, 0), 3, new AtomicInteger());
 		try {
 			template.setDefaultReplyTimeout(Duration.ofSeconds(5));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(E_REQUEST, null, null, null, "foo");
@@ -357,8 +358,9 @@ public void testAggregateTimeout() throws Exception {
 	@Test
 	@Disabled("time sensitive")
 	public void testAggregateTimeoutPartial() throws Exception {
+		AtomicInteger releaseCount = new AtomicInteger();
 		AggregatingReplyingKafkaTemplate<Integer, String, String> template = aggregatingTemplate(
-				new TopicPartitionOffset(F_REPLY, 0), 3);
+				new TopicPartitionOffset(F_REPLY, 0), 3, releaseCount);
 		template.setReturnPartialOnTimeout(true);
 		try {
 			template.setDefaultReplyTimeout(Duration.ofSeconds(5));
@@ -377,6 +379,7 @@ public void testAggregateTimeoutPartial() throws Exception {
 			assertThat(value2).isNotSameAs(value1);
 			assertThat(consumerRecord.topic())
 					.isEqualTo(AggregatingReplyingKafkaTemplate.PARTIAL_RESULTS_AFTER_TIMEOUT_TOPIC);
+			assertThat(releaseCount.get()).isEqualTo(2);
 		}
 		finally {
 			template.stop();
@@ -400,7 +403,8 @@ public void testAggregateOrphansNotStored() throws Exception {
 			correlation.set(rec.headers().lastHeader(KafkaHeaders.CORRELATION_ID).value());
 			return null;
 		}).given(producer).send(any(), any());
-		AggregatingReplyingKafkaTemplate template = new AggregatingReplyingKafkaTemplate(pf, container, coll -> true);
+		AggregatingReplyingKafkaTemplate template = new AggregatingReplyingKafkaTemplate(pf, container,
+				(list, timeout) -> true);
 		template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 		template.start();
 		List<ConsumerRecord> records = new ArrayList<>();
@@ -475,7 +479,7 @@ public ReplyingKafkaTemplate<Integer, String, String> createTemplate(TopicPartit
 	}
 
 	public AggregatingReplyingKafkaTemplate<Integer, String, String> aggregatingTemplate(
-			TopicPartitionOffset topic, int releaseSize) {
+			TopicPartitionOffset topic, int releaseSize, AtomicInteger releaseCount) {
 
 		ContainerProperties containerProperties = new ContainerProperties(topic);
 		containerProperties.setAckMode(AckMode.MANUAL_IMMEDIATE);
@@ -488,7 +492,11 @@ public AggregatingReplyingKafkaTemplate<Integer, String, String> aggregatingTemp
 				new KafkaMessageListenerContainer<>(cf, containerProperties);
 		container.setBeanName(this.testName);
 		AggregatingReplyingKafkaTemplate<Integer, String, String> template =
-				new AggregatingReplyingKafkaTemplate<>(this.config.pf(), container, coll -> coll.size() == releaseSize);
+				new AggregatingReplyingKafkaTemplate<>(this.config.pf(), container,
+						(list, timeout) -> {
+							releaseCount.incrementAndGet();
+							return list.size() == releaseSize;
+						});
 		template.setSharedReplyTopic(true);
 		template.start();
 		assertThat(template.getAssignedReplyTopicPartitions()).hasSize(1);

src/reference/asciidoc/kafka.adoc

@@ -503,11 +503,15 @@ The template in <<replying-template>> is strictly for a single request/reply sce
 For cases where multiple receivers of a single message return a reply, you can use the `AggregatingReplyingKafkaTemplate`.
 This is an implementation of the client-side of the https://www.enterpriseintegrationpatterns.com/patterns/messaging/BroadcastAggregate.html[Scatter-Gather Enterprise Integration Pattern].
 
-Like the `ReplyingKafkaTemplate`, the `AggregatingReplyingKafkaTemplate` constructor takes a producer factory and a listener container to receive the replies; it has a third parameter `Predicate<Collection<ConsumerRecord<K, R>>> releaseStrategy` which is consulted each time a reply is received; when the predicate returns `true`, the collection of `ConsumerRecord` s is used to complete the `Future` returned by the `sendAndReceive` method.
+Like the `ReplyingKafkaTemplate`, the `AggregatingReplyingKafkaTemplate` constructor takes a producer factory and a listener container to receive the replies; it has a third parameter `BiPredicate<List<ConsumerRecord<K, R>>, Boolean> releaseStrategy` which is consulted each time a reply is received; when the predicate returns `true`, the collection of `ConsumerRecord` s is used to complete the `Future` returned by the `sendAndReceive` method.
 
 There is an additional property `returnPartialOnTimeout` (default false).
 When this is set to `true`, instead of completing the future with a `KafkaReplyTimeoutException`, a partial result completes the future normally (as long as at least one reply record has been received).
 
+Starting with version 2.3.5, the predicate is also called after a timeout (if `returnPartialOnTimeout` is `true`).
+The first argument is the current list of records; the second is `true` if this call is due to a timeout.
+The predicate can modify the list of records.
+
 ====
 [source, java]
 ----

src/reference/asciidoc/whats-new.adoc

@@ -7,7 +7,7 @@ Also see <<new-in-sik>>.
 [[kafka-client-2.4]]
 ==== Kafka Client Version
 
-This version requires the 2.4.0 `kafka-clients` or higher.
+This version requires the 2.4.0 `kafka-clients` or higher and supports the new incremental rebalancing feature.
 
 [[x24-carl]]
 ==== ConsumerAwareRabalanceListener
@@ -25,6 +25,14 @@ See the IMPORTANT note at the end of <<rebalance-listeners>> for more informatio
 
 The `isAckAfterHandle()` default implementation now returns true by default.
 
+[[x24-agg]]
+==== AggregatingReplyingKafkaTemplate
+
+The `releastStrategy` is now a `BiConsumer`.
+It is now called after a timeout (as well as when records arrive); the second parameter is `true` in the case of a call after a timeout.
+
+See <<aggregating-request-reply>> for more information.
+
 === Migration Guide
 
 * This release is essentially the same as the 2.3.x line, except it has been compiled against the 2.4 `kafka-clients` jar, due to a binary incompatibility.