INT-2478 SimpleMessageStore and Sequence* Headers

Fix AbstractCorrelatingMessageHandler to ignore
Message Sequence information IF a custom release strategy is provided.

INT-2478 polishing based on PR comments.

Rename Test case.

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2011 the original author or authors.
+ * Copyright 2002-2012 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
@@ -35,6 +35,7 @@ import org.springframework.integration.store.MessageGroup;
 import org.springframework.integration.store.MessageGroupCallback;
 import org.springframework.integration.store.MessageGroupStore;
 import org.springframework.integration.store.MessageStore;
+import org.springframework.integration.store.SimpleMessageGroup;
 import org.springframework.integration.store.SimpleMessageStore;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -86,6 +87,8 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 
 	private final ConcurrentMap<Object, Object> locks = new ConcurrentHashMap<Object, Object>();
 
+	private volatile boolean sequenceAware = false;
+
 	public AbstractCorrelatingMessageHandler(MessageGroupProcessor processor, MessageGroupStore store,
 									 CorrelationStrategy correlationStrategy, ReleaseStrategy releaseStrategy) {
 		Assert.notNull(processor);
@@ -96,6 +99,7 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 				new HeaderAttributeCorrelationStrategy(MessageHeaders.CORRELATION_ID) : correlationStrategy;
 		this.releaseStrategy = releaseStrategy == null ? new SequenceSizeReleaseStrategy() : releaseStrategy;
 		this.messagingTemplate.setSendTimeout(DEFAULT_SEND_TIMEOUT);
+		sequenceAware = this.releaseStrategy instanceof SequenceSizeReleaseStrategy;
 	}
 
 	public AbstractCorrelatingMessageHandler(MessageGroupProcessor processor, MessageGroupStore store) {
@@ -124,6 +128,7 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 	public void setReleaseStrategy(ReleaseStrategy releaseStrategy) {
 		Assert.notNull(releaseStrategy);
 		this.releaseStrategy = releaseStrategy;
+		sequenceAware = this.releaseStrategy instanceof SequenceSizeReleaseStrategy;
 	}
 
 	public void setOutputChannel(MessageChannel outputChannel) {
@@ -182,16 +187,20 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 
 		synchronized (lock) {
 			MessageGroup messageGroup = messageStore.getMessageGroup(correlationKey);
+			if (this.sequenceAware){
+				messageGroup = new SequenceAwareMessageGroup(messageGroup);
+			}
+
 			if (!messageGroup.isComplete() && messageGroup.canAdd(message)) {
 				if (logger.isTraceEnabled()) {
 					logger.trace("Adding message to group [ " + messageGroup + "]");
 				}
-				messageGroup = store(correlationKey, message);
+				messageGroup = this.store(correlationKey, message);
 				
 				if (releaseStrategy.canRelease(messageGroup)) {
 					Collection<Message<?>> completedMessages = null;
 					try {
-						completedMessages = completeGroup(message, correlationKey, messageGroup);
+						completedMessages = this.completeGroup(message, correlationKey, messageGroup);
 					}
 					finally {
 						// Always clean up even if there was an exception
@@ -210,6 +219,7 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 		}
 	}
 
+
 	/**
 	 * Allows you to provide additional logic that needs to be performed after the MessageGroup was released. 
 	 * @param group
@@ -363,4 +373,43 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 		return false;
 	}
 
+	private static class SequenceAwareMessageGroup extends SimpleMessageGroup {
+
+		public SequenceAwareMessageGroup(MessageGroup messageGroup) {
+			super(messageGroup);
+		}
+
+		/**
+		 * This method determines whether messages have been added to this group that supersede the given message based on
+		 * its sequence id. This can be helpful to avoid ending up with sequences larger than their required sequence size
+		 * or sequences that are missing certain sequence numbers.
+		 */
+		public boolean canAdd(Message<?> message) {
+			if (this.size() == 0) {
+				return true;
+			}
+			Integer messageSequenceNumber = message.getHeaders().getSequenceNumber();
+			if (messageSequenceNumber != null && messageSequenceNumber > 0) {
+				Integer messageSequenceSize = message.getHeaders().getSequenceSize();
+				if (!messageSequenceSize.equals(this.getSequenceSize())) {
+					return false;
+				}
+				else {
+					return !this.containsSequenceNumber(this.getMessages(), messageSequenceNumber);
+				}
+			}
+			return true;
+		}
+
+		private boolean containsSequenceNumber(Collection<Message<?>> messages, Integer messageSequenceNumber) {
+			for (Message<?> member : messages) {
+				Integer memberSequenceNumber = member.getHeaders().getSequenceNumber();
+				if (messageSequenceNumber.equals(memberSequenceNumber)) {
+					return true;
+				}
+			}
+			return false;
+		}
+	}
+
 }

spring-integration-core/src/main/java/org/springframework/integration/store/MessageGroup.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2011 the original author or authors.
+ * Copyright 2002-2012 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.

spring-integration-core/src/main/java/org/springframework/integration/store/SimpleMessageGroup.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2011 the original author or authors.
+ * Copyright 2002-2012 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
@@ -77,7 +77,7 @@ public class SimpleMessageGroup implements MessageGroup {
 	}
 
 	public boolean canAdd(Message<?> message) {
-		return !isMember(message);
+		return true;
 	}
 
 	public void add(Message<?> message) {
@@ -136,38 +136,6 @@ public class SimpleMessageGroup implements MessageGroup {
 		this.messages.clear();
 	}
 
-	/**
-	 * This method determines whether messages have been added to this group that supersede the given message based on
-	 * its sequence id. This can be helpful to avoid ending up with sequences larger than their required sequence size
-	 * or sequences that are missing certain sequence numbers.
-	 */
-	private boolean isMember(Message<?> message) {
-		if (size() == 0) {
-			return false;
-		}
-		Integer messageSequenceNumber = message.getHeaders().getSequenceNumber();
-		if (messageSequenceNumber != null && messageSequenceNumber > 0) {
-			Integer messageSequenceSize = message.getHeaders().getSequenceSize();
-			if (!messageSequenceSize.equals(getSequenceSize())) {
-				return true;
-			}
-			else {
-				return this.containsSequenceNumber(messages, messageSequenceNumber);
-			}
-		}
-		return false;
-	}
-
-	private boolean containsSequenceNumber(Collection<Message<?>> messages, Integer messageSequenceNumber) {
-		for (Message<?> member : messages) {
-			Integer memberSequenceNumber = member.getHeaders().getSequenceNumber();
-			if (messageSequenceNumber.equals(memberSequenceNumber)) {
-				return true;
-			}
-		}
-		return false;
-	}
-
 	@Override
 	public String toString() {
 		return "SimpleMessageGroup{" +

spring-integration-core/src/test/java/org/springframework/integration/aggregator/scenarios/AggregatorWithCustomReleaseStrategyTests.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2002-2012 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
+ *
+ * http://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.integration.aggregator.scenarios;
+
+import org.junit.Test;
+
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.support.ClassPathXmlApplicationContext;
+import org.springframework.integration.Message;
+import org.springframework.integration.MessageChannel;
+import org.springframework.integration.channel.QueueChannel;
+import org.springframework.integration.support.MessageBuilder;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author Oleg Zhurakousky
+ *
+ */
+public class AggregatorWithCustomReleaseStrategyTests {
+
+	@Test
+	public void validateSequenceSizeHasNoAffect() throws Exception{
+		ApplicationContext context =
+				new ClassPathXmlApplicationContext("aggregator-with-custom-release-strategy.xml", this.getClass());
+		final MessageChannel inputChannel = context.getBean("in", MessageChannel.class);
+		QueueChannel resultChannel = context.getBean("resultChannel", QueueChannel.class);
+
+		new Thread(new Runnable() {
+			public void run() {
+				inputChannel.send(MessageBuilder.withPayload(new Integer[]{1, 2, 3, 4, 5, 6, 7, 8}).
+						setHeader("correlation", "foo").build());
+			}
+		}).start();
+
+		new Thread(new Runnable() {
+			public void run() {
+				inputChannel.send(MessageBuilder.withPayload(new Integer[]{10, 20, 30, 40, 50, 60, 70, 80}).
+						setHeader("correlation", "foo").build());
+			}
+		}).start();
+
+
+		Message<?> message = resultChannel.receive(1000);
+		int counter = 0;
+		while(message != null){
+			counter++;
+			message = resultChannel.receive(1000);
+		}
+		assertEquals(8, counter);
+	}
+}

spring-integration-core/src/test/java/org/springframework/integration/aggregator/scenarios/aggregator-with-custom-release-strategy.xml

@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<beans xmlns="http://www.springframework.org/schema/beans"
+	xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	xmlns:int="http://www.springframework.org/schema/integration"
+	xsi:schemaLocation="http://www.springframework.org/schema/integration http://www.springframework.org/schema/integration/spring-integration-2.1.xsd
+		http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans.xsd">
+
+
+	<int:splitter input-channel="in" output-channel="aggregationChannel"/>
+
+	<int:aggregator input-channel="aggregationChannel" output-channel="resultChannel"
+			release-strategy-expression="size() == 2"
+			correlation-strategy-expression="headers.correlation"
+			expire-groups-upon-completion="true"/>
+
+	<int:channel id="resultChannel">
+		<int:queue/>
+	</int:channel>
+</beans>

spring-integration-core/src/test/java/org/springframework/integration/store/SimpleMessageGroupTests.java

@@ -1,14 +1,17 @@
 package org.springframework.integration.store;
 
-import static org.hamcrest.CoreMatchers.is;
+import java.lang.reflect.Constructor;
-import static org.junit.Assert.assertThat;
-
 import java.util.Collections;
 
 import org.junit.Test;
+
 import org.springframework.integration.Message;
 import org.springframework.integration.support.MessageBuilder;
 
+import static org.hamcrest.CoreMatchers.is;
+
+import static org.junit.Assert.assertThat;
+
 /**
  * @author Iwein Fuld
  * @author Oleg Zhurakousky
@@ -20,8 +23,18 @@ public class SimpleMessageGroupTests {
 
 	private SimpleMessageGroup group = new SimpleMessageGroup(Collections.<Message<?>> emptyList(), key);
 
+	@SuppressWarnings("unchecked")
+	public void prepareForSequenceAwareMessageGroup() throws Exception{
+		Class<SimpleMessageGroup> clazz =
+				(Class<SimpleMessageGroup>)Class.forName("org.springframework.integration.aggregator.AbstractCorrelatingMessageHandler$SequenceAwareMessageGroup");
+		Constructor<SimpleMessageGroup> ctr = clazz.getDeclaredConstructor(MessageGroup.class);
+		ctr.setAccessible(true);
+		group = ctr.newInstance(group);
+	}
+
 	@Test
-	public void shouldFindSupersedingMessages() {
+	public void shouldFindSupersedingMessagesIfSequenceAware() throws Exception{
+		this.prepareForSequenceAwareMessageGroup();
 		final Message<?> message1 = MessageBuilder.withPayload("test").setSequenceNumber(1).build();
 		final Message<?> message2 = MessageBuilder.fromMessage(message1).setSequenceNumber(1).build();
 		assertThat(group.canAdd(message1), is(true));
@@ -31,7 +44,8 @@ public class SimpleMessageGroupTests {
 	}
 
 	@Test
-	public void shouldIgnoreMessagesWithZeroSequenceNumber() {
+	public void shouldIgnoreMessagesWithZeroSequenceNumberIfSequenceAware() throws Exception{
+		this.prepareForSequenceAwareMessageGroup();
 		final Message<?> message1 = MessageBuilder.withPayload("test").build();
 		final Message<?> message2 = MessageBuilder.fromMessage(message1).build();
 		assertThat(group.canAdd(message1), is(true));