Merge pull request #627 from garyrussell/INT-2751

* INT-2751:
  INT-2751 Fix Reaper Race Condition

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

@@ -244,20 +244,40 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 
 		Object correlationKey = group.getGroupId();
 		Lock lock = this.lockRegistry.obtain(correlationKey);
+		boolean removeGroup = true;
 		try {
 			lock.lockInterruptibly();
 			try {
 				if (group.size() > 0) {
 					try {
-						if (releaseStrategy.canRelease(group)) {
-							this.completeGroup(correlationKey, group);
+						/*
+						 * Need to verify the group hasn't changed while we were waiting on
+						 * its lock. We have to re-fetch the group for this. A possible
+						 * future improvement would be to add MessageGroupStore.getLastModified(groupId).
+						 */
+						MessageGroup messageGroupNow = this.messageStore.getMessageGroup(
+								group.getGroupId());
+						long lastModifiedNow = messageGroupNow.getLastModified();
+						if (group.getLastModified() == lastModifiedNow) {
+							if (releaseStrategy.canRelease(group)) {
+								this.completeGroup(correlationKey, group);
+							}
+							else {
+								this.expireGroup(correlationKey, group);
+							}
 						}
 						else {
-							this.expireGroup(correlationKey, group);
+							removeGroup = false;
+							if (logger.isDebugEnabled()) {
+								logger.debug("Group expiry candidate (" + group.getGroupId() +
+										") has changed - it may be reconsidered for a future expiration");
+							}
 						}
 					}
 					finally {
-						this.remove(group);
+						if (removeGroup) {
+							this.remove(group);
+						}
 					}
 					return true;
 				}

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

@@ -27,6 +27,7 @@ import org.springframework.integration.Message;
  * @author Iwein Fuld
  * @author Oleg Zhurakousky
  * @author Dave Syer
+ * @author Gary Russell
  * @since 2.0
  */
 public class SimpleMessageGroup implements MessageGroup {
@@ -144,6 +145,7 @@ public class SimpleMessageGroup implements MessageGroup {
 				"groupId=" + groupId +
 				", messages=" + messages +
 				", timestamp=" + timestamp +
+				", lastModified=" + lastModified +
 				'}';
 	}
 }

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

@@ -138,7 +138,9 @@ public class SimpleMessageStore extends AbstractMessageGroupStore implements Mes
 		if (group == null) {
 			return new SimpleMessageGroup(groupId);
 		}
-		return new SimpleMessageGroup(group);
+		SimpleMessageGroup simpleMessageGroup = new SimpleMessageGroup(group);
+		simpleMessageGroup.setLastModified(group.getLastModified());
+		return simpleMessageGroup;
 	}
 
 	public MessageGroup addMessageToGroup(Object groupId, Message<?> message) {
@@ -156,6 +158,7 @@ public class SimpleMessageStore extends AbstractMessageGroupStore implements Mes
 					this.groupIdToMessageGroup.putIfAbsent(groupId, group);
 				}
 				group.add(message);
+				this.groupIdToMessageGroup.get(groupId).setLastModified(System.currentTimeMillis());
 				return group;
 			}
 			finally {
@@ -199,6 +202,7 @@ public class SimpleMessageStore extends AbstractMessageGroupStore implements Mes
 				Assert.notNull(group, "MessageGroup for groupId '" + groupId + "' " +
 						"can not be located while attempting to remove Message from the MessageGroup");
 				group.remove(messageToRemove);
+				group.setLastModified(System.currentTimeMillis());
 				return group;
 			}
 			finally {
@@ -224,6 +228,7 @@ public class SimpleMessageStore extends AbstractMessageGroupStore implements Mes
 				Assert.notNull(group, "MessageGroup for groupId '" + groupId + "' " +
 						"can not be located while attempting to set 'lastReleasedSequenceNumber'");
 				group.setLastReleasedMessageSequenceNumber(sequenceNumber);
+				group.setLastModified(System.currentTimeMillis());
 			}
 			finally {
 				lock.unlock();
@@ -244,6 +249,7 @@ public class SimpleMessageStore extends AbstractMessageGroupStore implements Mes
 				Assert.notNull(group, "MessageGroup for groupId '" + groupId + "' " +
 						"can not be located while attempting to complete the MessageGroup");
 				group.complete();
+				group.setLastModified(System.currentTimeMillis());
 			}
 			finally {
 				lock.unlock();

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

@@ -0,0 +1,153 @@
+/*
+ * 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;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+import org.junit.Test;
+import org.springframework.integration.Message;
+import org.springframework.integration.MessageChannel;
+import org.springframework.integration.channel.QueueChannel;
+import org.springframework.integration.store.MessageGroup;
+import org.springframework.integration.store.MessageGroupStore;
+import org.springframework.integration.store.SimpleMessageStore;
+import org.springframework.integration.support.MessageBuilder;
+
+/**
+ * @author Gary Russell
+ * @since 2.2
+ *
+ */
+public class AbstractCorrelatingMessageHandlerTests {
+
+	@Test // INT-2751
+	public void testReaperDoesntReapAProcessingGroup() throws Exception {
+		final MessageGroupStore groupStore = new SimpleMessageStore();
+		final CountDownLatch waitForSendlatch = new CountDownLatch(1);
+		final CountDownLatch waitReapStartLatch = new CountDownLatch(1);
+		final CountDownLatch waitReapCompleteLatch = new CountDownLatch(1);
+		AbstractCorrelatingMessageHandler handler = new AbstractCorrelatingMessageHandler(
+				new MessageGroupProcessor() {
+
+					public Object processMessageGroup(MessageGroup group) {
+						return group;
+					}
+				}, groupStore) {
+
+			@Override
+			protected void afterRelease(MessageGroup group, Collection<Message<?>> completedMessages) {
+			}
+		};
+		handler.setReleasePartialSequences(true);
+
+		/*
+		 * Runs "reap" when group 'bar' is in completion
+		 */
+		Executors.newSingleThreadExecutor().execute(new Runnable() {
+
+			public void run() {
+				try {
+					waitReapStartLatch.await(10, TimeUnit.SECONDS);
+				}
+				catch (InterruptedException e) {
+					Thread.currentThread().interrupt();
+				}
+				waitForSendlatch.countDown();
+				try {
+					Thread.sleep(100);
+				}
+				catch (InterruptedException e) {
+					Thread.currentThread().interrupt();
+				}
+				groupStore.expireMessageGroups(50);
+				waitReapCompleteLatch.countDown();
+			}
+		});
+
+		final List<Message<?>> outputMessages = new ArrayList<Message<?>>();
+		handler.setOutputChannel(new MessageChannel() {
+
+			/*
+			 * Executes when group 'bar' completes normally
+			 */
+			public boolean send(Message<?> message, long timeout) {
+				outputMessages.add(message);
+				// wake reaper
+				waitReapStartLatch.countDown();
+				try {
+					waitForSendlatch.await(10, TimeUnit.SECONDS);
+					// wait a little longer for reaper to grab groups
+					Thread.sleep(2000);
+					// simulate tx commit
+					groupStore.removeMessageGroup("bar");
+				}
+				catch (InterruptedException e) {
+					Thread.currentThread().interrupt();
+				}
+				return true;
+			}
+
+			public boolean send(Message<?> message) {
+				return this.send(message, 0);
+			}
+		});
+		handler.setReleaseStrategy(new ReleaseStrategy() {
+
+			public boolean canRelease(MessageGroup group) {
+				return group.size() == 2;
+			}
+		});
+
+		QueueChannel discards = new QueueChannel();
+		handler.setDiscardChannel(discards);
+		handler.setSendPartialResultOnExpiry(true);
+
+		Message<String>	message = MessageBuilder.withPayload("foo")
+				.setCorrelationId("qux")
+				.build();
+		// partial group that will be reaped
+		handler.handleMessage(message);
+		message = MessageBuilder.withPayload("foo")
+				.setCorrelationId("bar")
+				.build();
+		// full group that should not be reaped
+		handler.handleMessage(message);
+		message = MessageBuilder.withPayload("baz")
+				.setCorrelationId("bar")
+				.build();
+		handler.handleMessage(message);
+
+		assertTrue(waitReapCompleteLatch.await(10, TimeUnit.SECONDS));
+		// Before INT-2751 we got bar + bar + qux
+		assertEquals(2, outputMessages.size()); // bar + qux
+		// normal release
+		assertEquals(2, ((MessageGroup) outputMessages.get(0).getPayload()).size()); // 'bar'
+		// reaper release
+		assertEquals(1, ((MessageGroup) outputMessages.get(1).getPayload()).size()); // 'qux'
+
+		assertNull(discards.receive(0));
+	}
+
+}