GH-3157: pop sequence for message as well

Fixes https://github.com/spring-projects/spring-integration/issues/3157

* Fixed issue with sequences not being popped when Output processor returns a message.

* Updated documents and added additional Checks to the testcase to ensure proper `popSequence` is done.

* Updated documents and added corrected code formatting

* Updated documents and used headerAccessor Constants for matching headers
* Clean up code style
* Improve docs for the feature on the matter

spring-integration-core/src/main/java/org/springframework/integration/aggregator/AbstractCorrelatingMessageHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -38,6 +38,7 @@ import org.springframework.context.Lifecycle;
 import org.springframework.expression.EvaluationContext;
 import org.springframework.expression.Expression;
 import org.springframework.integration.IntegrationMessageHeaderAccessor;
+import org.springframework.integration.StaticMessageHeaderAccessor;
 import org.springframework.integration.channel.NullChannel;
 import org.springframework.integration.context.IntegrationContextUtils;
 import org.springframework.integration.expression.ExpressionUtils;
@@ -59,6 +60,7 @@ import org.springframework.messaging.MessageHandlingException;
 import org.springframework.messaging.core.DestinationResolutionException;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
+import org.springframework.util.ObjectUtils;
 
 /**
  * Abstract Message handler that holds a buffer of correlated messages in a
@@ -89,6 +91,7 @@ import org.springframework.util.CollectionUtils;
  * @author David Liu
  * @author Enrique Rodriguez
  * @author Meherzad Lahewala
+ * @author Jayadev Sirimamilla
  *
  * @since 2.0
  */
@@ -535,43 +538,44 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageP
 	private void removeEmptyGroupAfterTimeout(MessageGroup messageGroup, long timeout) {
 		Object groupId = messageGroup.getGroupId();
 		UUID groupUuid = UUIDConverter.getUUID(groupId);
-		ScheduledFuture<?> scheduledFuture = getTaskScheduler()
-				.schedule(() -> {
-					Lock lock = this.lockRegistry.obtain(groupUuid.toString());
+		ScheduledFuture<?> scheduledFuture =
+				getTaskScheduler()
+						.schedule(() -> {
+							Lock lock = this.lockRegistry.obtain(groupUuid.toString());
 
-					try {
-						lock.lockInterruptibly();
-						try {
-							this.expireGroupScheduledFutures.remove(groupUuid);
-							/*
-							 * Obtain a fresh state for group from the MessageStore,
-							 * since it could be changed while we have waited for lock.
-							 */
-							MessageGroup groupNow = this.messageStore.getMessageGroup(groupUuid);
-							boolean removeGroup = groupNow.size() == 0 &&
-									groupNow.getLastModified()
-											<= (System.currentTimeMillis() - this.minimumTimeoutForEmptyGroups);
-							if (removeGroup) {
-								if (this.logger.isDebugEnabled()) {
-									this.logger.debug("Removing empty group: " + groupUuid);
+							try {
+								lock.lockInterruptibly();
+								try {
+									this.expireGroupScheduledFutures.remove(groupUuid);
+									/*
+									 * Obtain a fresh state for group from the MessageStore,
+									 * since it could be changed while we have waited for lock.
+									 */
+									MessageGroup groupNow = this.messageStore.getMessageGroup(groupUuid);
+									boolean removeGroup = groupNow.size() == 0 &&
+											groupNow.getLastModified()
+													<= (System.currentTimeMillis() - this.minimumTimeoutForEmptyGroups);
+									if (removeGroup) {
+										if (this.logger.isDebugEnabled()) {
+											this.logger.debug("Removing empty group: " + groupUuid);
+										}
+										remove(messageGroup);
+									}
+								}
+								finally {
+									lock.unlock();
 								}
-								remove(messageGroup);
 							}
-						}
-						finally {
-							lock.unlock();
-						}
-					}
-					catch (InterruptedException e) {
-						Thread.currentThread().interrupt();
-						if (this.logger.isDebugEnabled()) {
-							this.logger.debug("Thread was interrupted while trying to obtain lock."
-									+ "Rescheduling empty MessageGroup [ " + groupId + "] for removal.");
-						}
-						removeEmptyGroupAfterTimeout(messageGroup, timeout);
-					}
+							catch (InterruptedException e) {
+								Thread.currentThread().interrupt();
+								if (this.logger.isDebugEnabled()) {
+									this.logger.debug("Thread was interrupted while trying to obtain lock."
+											+ "Rescheduling empty MessageGroup [ " + groupId + "] for removal.");
+								}
+								removeEmptyGroupAfterTimeout(messageGroup, timeout);
+							}
 
-				}, new Date(System.currentTimeMillis() + timeout));
+						}, new Date(System.currentTimeMillis() + timeout));
 
 		if (this.logger.isDebugEnabled()) {
 			this.logger.debug("Schedule empty MessageGroup [ " + groupId + "] for removal.");
@@ -590,19 +594,20 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageP
 				final Object groupId = messageGroup.getGroupId();
 				final long timestamp = messageGroup.getTimestamp();
 				final long lastModified = messageGroup.getLastModified();
-				ScheduledFuture<?> scheduledFuture = getTaskScheduler()
-						.schedule(() -> {
-							try {
-								processForceRelease(groupId, timestamp, lastModified);
-							}
-							catch (MessageDeliveryException e) {
-								if (AbstractCorrelatingMessageHandler.this.logger.isWarnEnabled()) {
-									AbstractCorrelatingMessageHandler.this.logger.warn("The MessageGroup ["
-											+ groupId + "] is rescheduled by the reason of:", e);
-								}
-								scheduleGroupToForceComplete(groupId);
-							}
-						}, new Date(System.currentTimeMillis() + groupTimeout));
+				ScheduledFuture<?> scheduledFuture =
+						getTaskScheduler()
+								.schedule(() -> {
+									try {
+										processForceRelease(groupId, timestamp, lastModified);
+									}
+									catch (MessageDeliveryException e) {
+										if (AbstractCorrelatingMessageHandler.this.logger.isWarnEnabled()) {
+											AbstractCorrelatingMessageHandler.this.logger.warn("The MessageGroup ["
+													+ groupId + "] is rescheduled by the reason of:", e);
+										}
+										scheduleGroupToForceComplete(groupId);
+									}
+								}, new Date(System.currentTimeMillis() + groupTimeout));
 
 				if (this.logger.isDebugEnabled()) {
 					this.logger.debug("Schedule MessageGroup [ " + messageGroup + "] to 'forceComplete'.");
@@ -768,7 +773,7 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageP
 			Collection<Message<?>> partialSequence) {
 
 		Message<?> lastReleasedMessage = Collections.max(partialSequence, this.sequenceNumberComparator);
-		return new IntegrationMessageHeaderAccessor(lastReleasedMessage).getSequenceNumber();
+		return StaticMessageHeaderAccessor.getSequenceNumber(lastReleasedMessage);
 	}
 
 	protected MessageGroup store(Object correlationKey, Message<?> message) {
@@ -830,17 +835,23 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageP
 				partialSequence = (Collection<Message<?>>) result;
 			}
 
-			if (this.popSequence && partialSequence == null && !(result instanceof Message<?>)) {
-				AbstractIntegrationMessageBuilder<?> messageBuilder;
+			if (this.popSequence && partialSequence == null) {
+				AbstractIntegrationMessageBuilder<?> messageBuilder = null;
 				if (result instanceof AbstractIntegrationMessageBuilder<?>) {
 					messageBuilder = (AbstractIntegrationMessageBuilder<?>) result;
 				}
-				else {
-					messageBuilder = getMessageBuilderFactory()
+				else if (!(result instanceof Message<?>)) {
+					messageBuilder =
+							getMessageBuilderFactory()
 							.withPayload(result)
 							.copyHeaders(message.getHeaders());
 				}
-				result = messageBuilder.popSequenceDetails();
+				else if (compareSequences((Message<?>) result, message)) {
+					messageBuilder =
+							getMessageBuilderFactory()
+							.fromMessage((Message<?>) result);
+				}
+				result = messageBuilder != null ? messageBuilder.popSequenceDetails() : result;
 			}
 		}
 		finally {
@@ -852,6 +863,13 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageP
 		return partialSequence;
 	}
 
+	private static boolean compareSequences(Message<?> msg1, Message<?> msg2) {
+		Object sequence1 = msg1.getHeaders().get(IntegrationMessageHeaderAccessor.SEQUENCE_DETAILS);
+		Object sequence2 = msg2.getHeaders().get(IntegrationMessageHeaderAccessor.SEQUENCE_DETAILS);
+		return ObjectUtils.nullSafeEquals(sequence1, sequence2);
+
+	}
+
 	protected void verifyResultCollectionConsistsOfMessages(Collection<?> elements) {
 		Class<?> commonElementType = CollectionUtils.findCommonElementType(elements);
 		Assert.isAssignable(Message.class, commonElementType,

spring-integration-core/src/test/java/org/springframework/integration/dsl/routers/RouterTests.java

@@ -18,7 +18,6 @@ package org.springframework.integration.dsl.routers;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
-import static org.assertj.core.api.Assertions.fail;
 
 import java.util.Arrays;
 import java.util.List;
@@ -51,6 +50,7 @@ import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.MessageDeliveryException;
 import org.springframework.messaging.MessageHandlingException;
+import org.springframework.messaging.MessageHeaders;
 import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.PollableChannel;
 import org.springframework.messaging.core.DestinationResolutionException;
@@ -198,7 +198,6 @@ public class RouterTests {
 
 	@Test
 	public void testRecipientListRouter() {
-
 		Message<String> fooMessage = MessageBuilder.withPayload("fooPayload").setHeader("recipient", true).build();
 		Message<String> barMessage = MessageBuilder.withPayload("barPayload").setHeader("recipient", true).build();
 		Message<String> bazMessage = new GenericMessage<>("baz");
@@ -290,14 +289,9 @@ public class RouterTests {
 		assertThat(result2b).isNotNull();
 		assertThat(result2b.getPayload()).isEqualTo("bar");
 
-		try {
-			this.routerMethodInput.send(badMessage);
-			fail("MessageDeliveryException expected.");
-		}
-		catch (MessageDeliveryException e) {
-			assertThat(e.getMessage()).contains("No channel resolved by router");
-		}
-
+		assertThatExceptionOfType(MessageDeliveryException.class)
+				.isThrownBy(() -> this.routerMethodInput.send(badMessage))
+				.withMessageContaining("No channel resolved by router");
 	}
 
 	@Test
@@ -319,15 +313,10 @@ public class RouterTests {
 		assertThat(result2b).isNotNull();
 		assertThat(result2b.getPayload()).isEqualTo("bar");
 
-		try {
-			this.routerMethod2Input.send(badMessage);
-			fail("DestinationResolutionException expected.");
-		}
-		catch (MessagingException e) {
-			assertThat(e.getCause()).isInstanceOf(DestinationResolutionException.class);
-			assertThat(e.getCause().getMessage()).contains("failed to look up MessageChannel with name 'bad-channel'");
-		}
-
+		assertThatExceptionOfType(MessagingException.class)
+				.isThrownBy(() -> this.routerMethod2Input.send(badMessage))
+				.withCauseInstanceOf(DestinationResolutionException.class)
+				.withStackTraceContaining("failed to look up MessageChannel with name 'bad-channel'");
 	}
 
 	@Test
@@ -349,19 +338,14 @@ public class RouterTests {
 		assertThat(result2b).isNotNull();
 		assertThat(result2b.getPayload()).isEqualTo("bar");
 
-		try {
-			this.routerMethod3Input.send(badMessage);
-			fail("DestinationResolutionException expected.");
-		}
-		catch (MessagingException e) {
-			assertThat(e.getCause()).isInstanceOf(DestinationResolutionException.class);
-			assertThat(e.getCause().getMessage()).contains("failed to look up MessageChannel with name 'bad-channel'");
-		}
+		assertThatExceptionOfType(MessagingException.class)
+				.isThrownBy(() -> this.routerMethod3Input.send(badMessage))
+				.withCauseInstanceOf(DestinationResolutionException.class)
+				.withStackTraceContaining("failed to look up MessageChannel with name 'bad-channel'");
 	}
 
 	@Test
 	public void testMultiRouter() {
-
 		Message<String> fooMessage = new GenericMessage<>("foo");
 		Message<String> barMessage = new GenericMessage<>("bar");
 		Message<String> badMessage = new GenericMessage<>("bad");
@@ -382,13 +366,9 @@ public class RouterTests {
 		assertThat(result2b).isNotNull();
 		assertThat(result2b.getPayload()).isEqualTo("bar");
 
-		try {
-			this.routerMultiInput.send(badMessage);
-			fail("MessageDeliveryException expected.");
-		}
-		catch (MessageDeliveryException e) {
-			assertThat(e.getMessage()).contains("No channel resolved by router");
-		}
+		assertThatExceptionOfType(MessageDeliveryException.class)
+				.isThrownBy(() -> this.routerMultiInput.send(badMessage))
+				.withMessageContaining("No channel resolved by router");
 	}
 
 	@Autowired
@@ -612,6 +592,34 @@ public class RouterTests {
 
 	}
 
+	@Autowired
+	@Qualifier("scatterGatherWireTapChannel")
+	PollableChannel scatterGatherWireTapChannel;
+
+	@Test
+	public void testNestedScatterGatherSequenceTest() {
+		PollableChannel replyChannel = new QueueChannel();
+		this.scatterGatherInSubFlowChannel.send(
+				MessageBuilder.withPayload("sequencetest")
+						.setReplyChannel(replyChannel)
+						.build());
+
+		Message<?> wiretapMessage1 = scatterGatherWireTapChannel.receive(10000);
+		assertThat(wiretapMessage1).isNotNull();
+		MessageHeaders headers1 = wiretapMessage1.getHeaders();
+		Message<?> wiretapMessage2 = scatterGatherWireTapChannel.receive(10000);
+		assertThat(wiretapMessage2).isNotNull()
+				.extracting(Message::getHeaders)
+				.isEqualToComparingOnlyGivenFields(headers1, IntegrationMessageHeaderAccessor.CORRELATION_ID,
+						"gatherResultChannel", IntegrationMessageHeaderAccessor.SEQUENCE_NUMBER,
+						IntegrationMessageHeaderAccessor.SEQUENCE_SIZE);
+		Message<?> receive = replyChannel.receive(10000);
+
+		assertThat(receive).isNotNull();
+		assertThat(receive.getPayload()).isEqualTo("sequencetest");
+
+	}
+
 	@Configuration
 	@EnableIntegration
 	@EnableMessageHistory({ "recipientListOrder*", "recipient1*", "recipient2*" })
@@ -903,6 +911,7 @@ public class RouterTests {
 					.build();
 		}
 
+
 		@Bean
 		public IntegrationFlow propagateErrorFromGatherer(TaskExecutor taskExecutor) {
 			return IntegrationFlows.from(Function.class)
@@ -920,14 +929,21 @@ public class RouterTests {
 					.get();
 		}
 
+		@Bean
+		public PollableChannel scatterGatherWireTapChannel() {
+			return new QueueChannel();
+		}
+
 		@Bean
 		public IntegrationFlow scatterGatherInSubFlow() {
 			return flow -> flow.scatterGather(s -> s.applySequence(true)
-							.recipientFlow(inflow -> inflow
+							.recipientFlow(inflow -> inflow.wireTap(scatterGatherWireTapChannel())
 									.scatterGather(s1 -> s1.applySequence(true)
-													.recipientFlow(IntegrationFlowDefinition::bridge),
+													.recipientFlow(IntegrationFlowDefinition::bridge)
+													.recipientFlow("sequencetest"::equals,
+															IntegrationFlowDefinition::bridge),
 											g -> g.outputProcessor(MessageGroup::getOne)
-									)),
+									).wireTap(scatterGatherWireTapChannel()).bridge()),
 					g -> g.outputProcessor(MessageGroup::getOne));
 		}
 

src/reference/asciidoc/aggregator.adoc

@@ -114,9 +114,13 @@ This method is invoked for aggregating messages as follows:
 
 NOTE: In the interest of code simplicity and promoting best practices such as low coupling, testability, and others, the preferred way of implementing the aggregation logic is through a POJO and using the XML or annotation support for configuring it in the application.
 
-Starting with version 5.1, after processing message group, an `AbstractCorrelatingMessageHandler` performs a `MessageBuilder.popSequenceDetails()` message headers modification for the proper splitter-aggregator scenario with several nested levels.
-It is done only if the message group release result is not a message or collection of messages.
+Starting with version 5.3, after processing message group, an `AbstractCorrelatingMessageHandler` performs a `MessageBuilder.popSequenceDetails()` message headers modification for the proper splitter-aggregator scenario with several nested levels.
+It is done only if the message group release result is not a collection of messages.
 In that case a target `MessageGroupProcessor` is responsible for the `MessageBuilder.popSequenceDetails()` call while building those messages.
+
+If the `MessageGroupProcessor` returns a `Message`, a `MessageBuilder.popSequenceDetails()` will be performed on the output message only if the `sequenceDetails` matches with first message in group.
+(Previously this has been done only if a plain payload or an `AbstractIntegrationMessageBuilder` has been returned from the `MessageGroupProcessor`.)
+
 This functionality can be controlled by a new `popSequence` `boolean` property, so the `MessageBuilder.popSequenceDetails()` can be disabled in some scenarios when correlation details have not been populated by the standard splitter.
 This property, essentially, undoes what has been done by the nearest upstream `applySequence = true` in the `AbstractMessageSplitter`.
 See <<./splitter.adoc#splitter,Splitter>> for more information.
@@ -470,7 +474,7 @@ By default, an internal `DefaultLockRegistry` is used.
 Use of a distributed `LockRegistry`, such as the `ZookeeperLockRegistry`, ensures only one instance of the aggregator can operate on a group concurrently.
 See <<./redis.adoc#redis-lock-registry,Redis Lock Registry>>, <<./gemfire.adoc#gemfire-lock-registry,Gemfire Lock Registry>>, and <<./zookeeper.adoc#zk-lock-registry,Zookeeper Lock Registry>> for more information.
 <21> A timeout (in milliseconds) to force the `MessageGroup` complete when the `ReleaseStrategy` does not release the group when the current message arrives.
-This attribute provides a built-in time-based release strategy for the aggregator when there is a need to emit a partial result (or discard the group) if a new message does not arrive for the `MessageGroup` within the timeout which counts from the time the last message arrived. 
+This attribute provides a built-in time-based release strategy for the aggregator when there is a need to emit a partial result (or discard the group) if a new message does not arrive for the `MessageGroup` within the timeout which counts from the time the last message arrived.
 To set up a timeout which counts from the time the `MessageGroup` was created see `group-timeout-expression` information.
 When a new message arrives at the aggregator, any existing `ScheduledFuture<?>` for its `MessageGroup` is canceled.
 If the `ReleaseStrategy` returns `false` (meaning do not release) and `groupTimeout > 0`, a new task is scheduled to expire the group.
@@ -488,7 +492,7 @@ If it evaluates to zero, the group is completed immediately on the current threa
 In effect, this provides a dynamic `group-timeout` property.
 As an example, if you wish to forcibly complete a `MessageGroup` after 10 seconds have elapsed since the time the group was created you might consider using the following SpEL expression: `timestamp + 10000 - T(System).currentTimeMillis()` where `timestamp` is provided by `MessageGroup.getTimestamp()` as the `MessageGroup` here is the `#root` evaluation context object.
 Bear in mind however that the group creation time might differ from the time of the first arrived message depending on other group expiration properties' configuration.
-See `group-timeout` for more information. 
+See `group-timeout` for more information.
 Mutually exclusive with 'group-timeout' attribute.
 <23> When a group is completed due to a timeout (or by a `MessageGroupStoreReaper`), the group is expired (completely removed) by default.
 Late arriving messages start a new group.

src/reference/asciidoc/whats-new.adoc

@@ -33,6 +33,12 @@ See <<./reactive-streams.adoc/reactive-message-handler,ReactiveMessageHandler>>
 `spring-integration-mongodb` module now provides channel adapter implementations for Reactive MongoDB driver support in Spring Data.
 See <<./mongodb.adoc#mongodb-reactive-channel-adapters,MongoDB Reactive Channel Adapters>> for more information.
 
+[[x5.3-AbstractCorrelatingMessageHandler]]
+==== Aggregator Changes
+
+If the `MessageGroupProcessor` returns a `Message`, the `MessageBuilder.popSequenceDetails()` is performed on the output message if the `sequenceDetails` matches with first message of group.
+See <<./aggregator.adoc#aggregator-api,Aggregator Programming Model>> for more information.
+
 [[x5.3-general]]
 === General Changes