INT-4550: Disallow multi aggregators on same MGS (#2622)

* INT-4550: Disallow multi aggregators on same MGS

JIRA: https://jira.spring.io/browse/INT-4550

**Cherry-pick to 5.0.x**

* * Introduce `UniqueExpiryCallback`
* Use `UniqueExpiryCallback` in the `AbstractCorrelatingMessageHandler`
* Check for uniqueness in the `AbstractMessageGroupStore`
* Remove duplicate code in the `ConfigurableMongoDbMessageStore`

* * Fix tests according a new logic

* * Address PR review

* Change `Assert.isTrue` to the `logger.error` for backward compatibility
* Revert changes in tests since we don't throw exception anymore
* Fix language on doc

* * Fix Checkstyle violation in the `AbstractMessageGroupStore`

* * Ignore `testDontReapMessageOfOtherHandler()`

Author: artembilan

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

@@ -22,7 +22,6 @@
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledFuture;
@@ -51,6 +50,7 @@
 import org.springframework.integration.store.MessageStore;
 import org.springframework.integration.store.SimpleMessageGroup;
 import org.springframework.integration.store.SimpleMessageStore;
+import org.springframework.integration.store.UniqueExpiryCallback;
 import org.springframework.integration.support.locks.DefaultLockRegistry;
 import org.springframework.integration.support.locks.LockRegistry;
 import org.springframework.integration.util.UUIDConverter;
@@ -100,8 +100,6 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageP
 
 	private final Map<UUID, ScheduledFuture<?>> expireGroupScheduledFutures = new ConcurrentHashMap<>();
 
-	private final Set<Object> groupIds =  ConcurrentHashMap.newKeySet();
-
 	private MessageGroupProcessor outputProcessor;
 
 	private MessageGroupStore messageStore;
@@ -173,8 +171,9 @@ public void setLockRegistry(LockRegistry lockRegistry) {
 
 	public final void setMessageStore(MessageGroupStore store) {
 		this.messageStore = store;
-		store.registerMessageGroupExpiryCallback(
-				(messageGroupStore, group) -> this.forceReleaseProcessor.processMessageGroup(group));
+		UniqueExpiryCallback expiryCallback =
+				(messageGroupStore, group) -> this.forceReleaseProcessor.processMessageGroup(group);
+		store.registerMessageGroupExpiryCallback(expiryCallback);
 	}
 
 	public void setCorrelationStrategy(CorrelationStrategy correlationStrategy) {
@@ -687,7 +686,6 @@ protected void forceComplete(MessageGroup group) {
 	protected void remove(MessageGroup group) {
 		Object correlationKey = group.getGroupId();
 		this.messageStore.removeMessageGroup(correlationKey);
-		this.groupIds.remove(group.getGroupId());
 	}
 
 	protected int findLastReleasedSequenceNumber(Object groupId, Collection<Message<?>> partialSequence) {
@@ -696,7 +694,6 @@ protected int findLastReleasedSequenceNumber(Object groupId, Collection<Message<
 	}
 
 	protected MessageGroup store(Object correlationKey, Message<?> message) {
-		this.groupIds.add(correlationKey);
 		return this.messageStore.addMessageToGroup(correlationKey, message);
 	}
 
@@ -866,9 +863,7 @@ private class ForceReleaseMessageGroupProcessor implements MessageGroupProcessor
 
 		@Override
 		public Object processMessageGroup(MessageGroup group) {
-			if (AbstractCorrelatingMessageHandler.this.groupIds.contains(group.getGroupId())) {
-				forceComplete(group);
-			}
+			forceComplete(group);
 			return null;
 		}
 

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

@@ -42,7 +42,7 @@ public abstract class AbstractMessageGroupStore extends AbstractBatchingMessageG
 
 	protected final Log logger = LogFactory.getLog(getClass());
 
-	private final Collection<MessageGroupCallback> expiryCallbacks = new LinkedHashSet<MessageGroupCallback>();
+	private final Collection<MessageGroupCallback> expiryCallbacks = new LinkedHashSet<>();
 
 	private final MessageGroupFactory persistentMessageGroupFactory =
 			new SimpleMessageGroupFactory(SimpleMessageGroupFactory.GroupType.PERSISTENT);
@@ -72,13 +72,10 @@ protected MessageGroupFactory getMessageGroupFactory() {
 	/**
 	 * Convenient injection point for expiry callbacks in the message store. Each of the callbacks provided will simply
 	 * be registered with the store using {@link #registerMessageGroupExpiryCallback(MessageGroupCallback)}.
-	 *
 	 * @param expiryCallbacks the expiry callbacks to add
 	 */
 	public void setExpiryCallbacks(Collection<MessageGroupCallback> expiryCallbacks) {
-		for (MessageGroupCallback callback : expiryCallbacks) {
-			registerMessageGroupExpiryCallback(callback);
-		}
+		expiryCallbacks.forEach(this::registerMessageGroupExpiryCallback);
 	}
 
 	public boolean isTimeoutOnIdle() {
@@ -110,6 +107,17 @@ public void setLazyLoadMessageGroups(boolean lazyLoadMessageGroups) {
 
 	@Override
 	public void registerMessageGroupExpiryCallback(MessageGroupCallback callback) {
+		if (callback instanceof UniqueExpiryCallback) {
+			boolean uniqueExpiryCallbackPresent =
+					this.expiryCallbacks.stream()
+							.anyMatch(UniqueExpiryCallback.class::isInstance);
+
+			if (!uniqueExpiryCallbackPresent && this.logger.isErrorEnabled()) {
+				this.logger.error("Only one instance of 'UniqueExpiryCallback' can be registered in the " +
+						"'MessageGroupStore'. Use a separate 'MessageGroupStore' for each aggregator/resequencer.");
+			}
+		}
+
 		this.expiryCallbacks.add(callback);
 	}
 

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

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2018 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.store;
+
+/**
+ * A marker interface extension of the {@link MessageGroupStore.MessageGroupCallback}
+ * for components which should be registered in the {@link MessageGroupStore} only once.
+ * The {@link MessageGroupStore} implementation ensures that only once instance of this
+ * class is present in the expire callbacks.
+ *
+ * @author Meherzad Lahewala
+ * @author Artme Bilan
+ *
+ * @since 5.0.10
+ */
+@FunctionalInterface
+public interface UniqueExpiryCallback extends MessageGroupStore.MessageGroupCallback {
+
+}

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

@@ -35,6 +35,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 
+import org.junit.Ignore;
 import org.junit.Test;
 
 import org.springframework.beans.DirectFieldAccessor;
@@ -422,7 +423,8 @@ public void testScheduleRemoveAnEmptyGroupAfterConfiguredDelay() throws Exceptio
 	}
 
 	@Test
-	public void testDontReapMessageOfOtherHandler() throws Exception {
+	@Ignore("Until 5.2 with new 'owner' feature on groups")
+	public void testDontReapMessageOfOtherHandler() {
 		MessageGroupStore groupStore = new SimpleMessageStore();
 
 		AggregatingMessageHandler handler1 = new AggregatingMessageHandler(group -> group, groupStore);

spring-integration-mongodb/src/main/java/org/springframework/integration/mongodb/store/ConfigurableMongoDbMessageStore.java

@@ -17,12 +17,12 @@
 package org.springframework.integration.mongodb.store;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.Iterator;
-import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.UUID;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
 
 import org.springframework.data.domain.Sort;
 import org.springframework.data.mongodb.MongoDbFactory;
@@ -32,12 +32,10 @@
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
 import org.springframework.integration.store.MessageGroup;
-import org.springframework.integration.store.MessageGroupMetadata;
 import org.springframework.integration.store.MessageGroupStore;
 import org.springframework.integration.store.MessageStore;
 import org.springframework.integration.store.SimpleMessageGroup;
 import org.springframework.jmx.export.annotation.ManagedAttribute;
-import org.springframework.jmx.export.annotation.ManagedOperation;
 import org.springframework.messaging.Message;
 import org.springframework.util.Assert;
 
@@ -58,11 +56,6 @@ public class ConfigurableMongoDbMessageStore extends AbstractConfigurableMongoDb
 
 	public final static String DEFAULT_COLLECTION_NAME = "configurableStoreMessages";
 
-	private final Collection<MessageGroupCallback> expiryCallbacks = new LinkedHashSet<MessageGroupCallback>();
-
-	private volatile boolean timeoutOnIdle;
-
-
 	public ConfigurableMongoDbMessageStore(MongoTemplate mongoTemplate) {
 		this(mongoTemplate, DEFAULT_COLLECTION_NAME);
 	}
@@ -85,33 +78,8 @@ public ConfigurableMongoDbMessageStore(MongoDbFactory mongoDbFactory, String col
 
 	public ConfigurableMongoDbMessageStore(MongoDbFactory mongoDbFactory, MappingMongoConverter mappingMongoConverter,
 			String collectionName) {
-		super(mongoDbFactory, mappingMongoConverter, collectionName);
-	}
-
-	/**
-	 * Convenient injection point for expiry callbacks in the message store. Each of the callbacks provided will simply
-	 * be registered with the store using {@link #registerMessageGroupExpiryCallback(MessageGroupCallback)}.
-	 * @param expiryCallbacks the expiry callbacks to add
-	 */
-	public void setExpiryCallbacks(Collection<MessageGroupCallback> expiryCallbacks) {
-		for (MessageGroupCallback callback : expiryCallbacks) {
-			registerMessageGroupExpiryCallback(callback);
-		}
-	}
-
-	public boolean isTimeoutOnIdle() {
-		return this.timeoutOnIdle;
-	}
 
-	/**
-	 * Allows you to override the rule for the timeout calculation. Typical timeout is based from the time
-	 * the {@link MessageGroup} was created. If you want the timeout to be based on the time
-	 * the {@link MessageGroup} was idling (e.g., inactive from the last update) invoke this method with 'true'.
-	 * Default is 'false'.
-	 * @param timeoutOnIdle The boolean.
-	 */
-	public void setTimeoutOnIdle(boolean timeoutOnIdle) {
-		this.timeoutOnIdle = timeoutOnIdle;
+		super(mongoDbFactory, mappingMongoConverter, collectionName);
 	}
 
 	@Override
@@ -204,7 +172,7 @@ public void removeMessagesFromGroup(Object groupId, Collection<Message<?>> messa
 		Assert.notNull(groupId, "'groupId' must not be null");
 		Assert.notNull(messages, "'messageToRemove' must not be null");
 
-		Collection<UUID> ids = new ArrayList<UUID>();
+		Collection<UUID> ids = new ArrayList<>();
 		for (Message<?> messageToRemove : messages) {
 			ids.add(messageToRemove.getHeaders().getId());
 			if (ids.size() >= getRemoveBatchSize()) {
@@ -224,11 +192,6 @@ private void removeMessages(Object groupId, Collection<UUID> ids) {
 		this.mongoTemplate.remove(query, this.collectionName);
 	}
 
-	@Override
-	public void removeMessagesFromGroup(Object groupId, Message<?>... messages) {
-		removeMessagesFromGroup(groupId, Arrays.asList(messages));
-	}
-
 	@Override
 	public Message<?> pollMessageFromGroup(final Object groupId) {
 		Assert.notNull(groupId, "'groupId' must not be null");
@@ -246,52 +209,24 @@ public Message<?> pollMessageFromGroup(final Object groupId) {
 
 	@Override
 	public void setLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber) {
-		this.updateGroup(groupId, lastModifiedUpdate().set(MessageDocumentFields.LAST_RELEASED_SEQUENCE, sequenceNumber));
+		updateGroup(groupId, lastModifiedUpdate().set(MessageDocumentFields.LAST_RELEASED_SEQUENCE, sequenceNumber));
 	}
 
 	@Override
 	public void completeGroup(Object groupId) {
-		this.updateGroup(groupId, lastModifiedUpdate().set(MessageDocumentFields.COMPLETE, true));
+		updateGroup(groupId, lastModifiedUpdate().set(MessageDocumentFields.COMPLETE, true));
 	}
 
 	@Override
 	public Iterator<MessageGroup> iterator() {
-		List<MessageGroup> messageGroups = new ArrayList<MessageGroup>();
-
 		Query query = Query.query(Criteria.where(MessageDocumentFields.GROUP_ID).exists(true));
 		Iterable<String> groupIds = mongoTemplate.getCollection(collectionName)
 				.distinct(MessageDocumentFields.GROUP_ID, query.getQueryObject(), String.class);
 
-		for (Object groupId : groupIds) {
-			messageGroups.add(getMessageGroup(groupId));
-		}
-
-		return messageGroups.iterator();
-	}
-
-	@Override
-	public void registerMessageGroupExpiryCallback(MessageGroupCallback callback) {
-		this.expiryCallbacks.add(callback);
-	}
-
-	@Override
-	@ManagedOperation
-	public int expireMessageGroups(long timeout) {
-		int count = 0;
-		long threshold = System.currentTimeMillis() - timeout;
-		for (MessageGroup group : this) {
-
-			long timestamp = group.getTimestamp();
-			if (this.isTimeoutOnIdle() && group.getLastModified() > 0) {
-				timestamp = group.getLastModified();
-			}
+		return StreamSupport.stream(groupIds.spliterator(), false)
+				.map(this::getMessageGroup)
+				.iterator();
 
-			if (timestamp <= threshold) {
-				count++;
-				expire(group);
-			}
-		}
-		return count;
 	}
 
 	@Override
@@ -314,11 +249,6 @@ public int getMessageGroupCount() {
 				.size();
 	}
 
-	@Override
-	public MessageGroupMetadata getGroupMetadata(Object groupId) {
-		throw new UnsupportedOperationException("Not yet implemented for this store");
-	}
-
 	@Override
 	public Message<?> getOneMessageFromGroup(Object groupId) {
 		Assert.notNull(groupId, "'groupId' must not be null");
@@ -337,36 +267,12 @@ public Collection<Message<?>> getMessagesForGroup(Object groupId) {
 		Assert.notNull(groupId, "'groupId' must not be null");
 		Query query = groupOrderQuery(groupId);
 		List<MessageDocument> documents = this.mongoTemplate.find(query, MessageDocument.class, this.collectionName);
-		List<Message<?>> messages = new ArrayList<Message<?>>();
-
-		for (MessageDocument document : documents) {
-			messages.add(document.getMessage());
-		}
-		return messages;
-	}
-
-	private void expire(MessageGroup group) {
-
-		RuntimeException exception = null;
-
-		for (MessageGroupCallback callback : this.expiryCallbacks) {
-			try {
-				callback.execute(this, group);
-			}
-			catch (RuntimeException e) {
-				if (exception == null) {
-					exception = e;
-				}
-				logger.error("Exception in expiry callback", e);
-			}
-		}
 
-		if (exception != null) {
-			throw exception;
-		}
+		return documents.stream()
+				.map(MessageDocument::getMessage)
+				.collect(Collectors.toList());
 	}
 
-
 	private void updateGroup(Object groupId, Update update) {
 		this.mongoTemplate.updateFirst(groupOrderQuery(groupId), update, this.collectionName);
 	}

src/reference/asciidoc/aggregator.adoc

@@ -814,9 +814,16 @@ by removing the group from the store entirely).
 
 The `MessageGroupStore` maintains a list of these callbacks which it applies, on demand, to all messages whose timestamp is earlier than a time supplied as a parameter (see the `registerMessageGroupExpiryCallback(..)` and `expireMessageGroups(..)` methods above).
 
-The `expireMessageGroups` method can be called with a timeout value.
-Any message older than the current time minus this value will be expired, and have the callbacks applied.
-Thus it is the user of the store that defines what is meant by message group "expiry".
+IMPORTANT: It is important not to use the same `MessageGroupStore` instance in different aggregator components, when you intend to rely on the `expireMessageGroups` functionality.
+Every `AbstractCorrelatingMessageHandler` registers its own `MessageGroupCallback` based on the `forceComplete()` callback.
+This way each group for expiration may be completed or discarded by the wrong aggregator.
+Starting with version 5.0.10, a `UniqueExpiryCallback` is used from the `AbstractCorrelatingMessageHandler` for the registration callback in the `MessageGroupStore`.
+The `MessageGroupStore`, in turn, checks for presence an instance of this class and logs an error with an appropriate message if one is already present in the callbacks set.
+This way the Framework disallows usage of the `MessageGroupStore` instance in different aggregators/resequencers to avoid the mentioned side effect of expiration the groups not created by the particular correlation handler.
+
+You can call the `expireMessageGroups` method with a timeout value.
+Any message older than the current time minus this value is expired and has the callbacks applied.
+Thus, it is the user of the store that defines what is meant by message group "`expiry`".
 
 As a convenience for users, Spring Integration provides a wrapper for the message expiry in the form of a `MessageGroupStoreReaper`: