OPEN - issue INT-1063: MessageStore: correlation and grouping API

Here is something that works, but there is more work to do in simplifying the CorrelatingMessageHandler and friends

org.springframework.integration/src/main/java/org/springframework/integration/aggregator/CorrelatingMessageHandler.java

@@ -186,7 +186,7 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 		try {
 			if (tracker.waitForLockIfNotTracked(correlationKey)) {
 				MessageGroup group = new MessageGroup(store.list(correlationKey),
-						completionStrategy, correlationKey, deleteOrTrackCallback());
+						completionStrategy, correlationKey, deleteOrTrackCallback(correlationKey));
 
 				if (group.hasNoMessageSuperseding(message)) {
 					store(message, correlationKey);
@@ -212,17 +212,18 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 		}
 	}
 
-	private MessageGroupListener deleteOrTrackCallback() {
+	private MessageGroupListener deleteOrTrackCallback(final Object correlationKey) {
 		return new MessageGroupListener() {
-
+			
 			public void onProcessingOf(Message<?>... processedMessage) {
 				for (Message<?> message : processedMessage) {
-					store.delete(message.getHeaders().getId());
+					store.delete(correlationKey, message.getHeaders().getId());
 				}
 			}
 
 			public void onCompletionOf(Object correlationKey) {
 				tracker.pushCorrelationId(correlationKey);
+				store.deleteAll(correlationKey);
 			}
 		};
 	}
@@ -233,7 +234,7 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 		if (!correlationKey.equals(message.getHeaders().getCorrelationId())) {
 			toStore = MessageBuilder.fromMessage(message).setCorrelationId(correlationKey).build();
 		}
-		store.put(toStore);
+		store.put(correlationKey, toStore);
 		if (!keysInBuffer.contains(correlationKey)) {
 			keysInBuffer.add(new DelayedKey(correlationKey, timeout));
 		}
@@ -293,7 +294,7 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 		try {
 			if (tracker.tryLockFor(key)) {
 				Collection<Message<?>> all = store.list(key);
-				MessageGroup group = new MessageGroup(all, completionStrategy, key, deleteOrTrackCallback());
+				MessageGroup group = new MessageGroup(all, completionStrategy, key, deleteOrTrackCallback(key));
 				if (all.size() > 0) {
 					// last chance for normal completion
 					MessageChannel outputChannel = resolveReplyChannel(all.iterator().next(), this.outputChannel);

org.springframework.integration/src/main/java/org/springframework/integration/store/MessageStore.java

@@ -59,10 +59,49 @@ public interface MessageStore {
 	Message<?> delete(UUID id);
 
 	/**
-	 * Return all Messages currently in the MessageStore that contain the
-	 * provided correlationId header value.
+	 * Return all Messages currently in the MessageStore that were stored using
+	 * {@link #put(Object, Message)} or {@link #put(Object, Collection)} with
+	 * this correlation id.
+	 * 
 	 * @see org.springframework.integration.core.MessageHeaders#getCorrelationId()
 	 */
 	Collection<Message<?>> list(Object correlationId);
 
+	/**
+	 * Store a message with an association to a correlation id. This can be used
+	 * to group messages together instead of storing them just under their id.
+	 * 
+	 * @param correlationId the correlation id to store the message under
+	 * @param message a message
+	 */
+	void put(Object correlationId, Message<?> message);
+
+	/**
+	 * Store a group of message with an association to a correlation id.
+	 * 
+	 * @param correlationId the correlation id to store the message under
+	 * @param messages a collection of messages
+	 * 
+	 * @see MessageStore#put(UUID, Message)
+	 */
+	void put(Object correlationId, Collection<Message<?>> messages);
+
+	/**
+	 * Delete a message from the association with this correlation id. If the
+	 * message was stored under through {@link #put(Message)} as well, then it
+	 * is still accessible via {@link #get(UUID)}.
+	 * 
+	 * @param correlationId the correlation id to delete all messages under
+	 */
+	Message<?> delete(Object correlationId, UUID messageId);
+
+	/**
+	 * Delete all the messages from the association with this correlation id. If
+	 * the messages were stored under their id through {@link #put(Message)}
+	 * they are still accessible via {@link #get(UUID)}.
+	 * 
+	 * @param correlationId the correlation id to delete all messages under
+	 */
+	void deleteAll(Object correlationId);
+
 }

org.springframework.integration/src/main/java/org/springframework/integration/store/SimpleMessageStore.java

@@ -16,12 +16,13 @@
 
 package org.springframework.integration.store;
 
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.List;
-import java.util.Map;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Iterator;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 import org.springframework.integration.core.Message;
 import org.springframework.integration.core.MessagingException;
@@ -38,7 +39,9 @@ import org.springframework.util.Assert;
  */
 public class SimpleMessageStore implements MessageStore {
 
-	private final Map<UUID, Message<?>> map;
+	private final ConcurrentMap<UUID, Message<?>> idToMessage;
+
+	private final ConcurrentMap<Object, Collection<Message<?>>> correlationToMessage;
 
 	private final UpperBound upperBound;
 
@@ -47,7 +50,8 @@ public class SimpleMessageStore implements MessageStore {
 	 * capacity, or unlimited size if the given capacity is less than 1.
 	 */
 	public SimpleMessageStore(int capacity) {
-		this.map = new ConcurrentHashMap<UUID, Message<?>>();
+		this.idToMessage = new ConcurrentHashMap<UUID, Message<?>>();
+		this.correlationToMessage = new ConcurrentHashMap<Object, Collection<Message<?>>>();
 		this.upperBound = new UpperBound(capacity);
 	}
 
@@ -64,37 +68,73 @@ public class SimpleMessageStore implements MessageStore {
 			throw new MessagingException(this.getClass().getSimpleName()
 					+ " was out of capacity at, try constructing it with a larger capacity.");
 		}
-		return (Message<T>) this.map.put(message.getHeaders().getId(), message);
+		Object correlationId = message.getHeaders().getCorrelationId();
+		if (correlationId!=null) {
+			getMessagesInternal(correlationId).add(message);
+		}
+		return (Message<T>) this.idToMessage.put(message.getHeaders().getId(), message);
 	}
 
 	public Message<?> get(UUID key) {
-		return (key != null) ? this.map.get(key) : null;
+		return (key != null) ? this.idToMessage.get(key) : null;
 	}
 
 	public Message<?> delete(UUID key) {
 		if (key != null) {
 			upperBound.release();
-			return this.map.remove(key);
+			return this.idToMessage.remove(key);
 		}
 		else
 			return null;
 	}
 
 	public int size() {
-		return this.map.size();
+		return this.idToMessage.size();
 	}
 
-	public List<Message<?>> list(Object correlationKey) {
-		Assert.notNull(correlationKey, "'correlationKey' must not be null");
-		List<Message<?>> matched = new ArrayList<Message<?>>();
-		Collection<Message<?>> values = map.values();
-		for (Message<?> message : values) {
-			Object correlationId = message.getHeaders().getCorrelationId();
-			if (correlationId != null && correlationId.equals(correlationKey)) {
-				matched.add(message);
+	public Collection<Message<?>> list(Object correlationId) {
+		Assert.notNull(correlationId, "'correlationKey' must not be null");
+		Collection<Message<?>> collection = correlationToMessage.get(correlationId);
+		if (collection==null) {
+			return Collections.emptySet();
+		}
+		return Collections.unmodifiableCollection(collection);
+	}
+
+	public void put(Object correlationId, Collection<Message<?>> messages) {
+		getMessagesInternal(correlationId).addAll(messages);
+	}
+
+	public void put(Object correlationId, Message<?> message) {
+		getMessagesInternal(correlationId).add(message);
+	}
+
+	public Message<?> delete(Object correlationId, UUID messageId) {
+		if (!correlationToMessage.containsKey(correlationId)) {
+			return null;
+		}
+		Collection<Message<?>> messages = getMessagesInternal(correlationId);
+		Message<?> result = null;
+		for (Iterator<Message<?>> iterator = messages.iterator(); iterator.hasNext();) {
+			Message<?> message = (Message<?>) iterator.next();
+			if (message.getHeaders().getId().equals(messageId)) {
+				iterator.remove();
+				result = message;
 			}
 		}
-		return matched;
+		return result;
+	}
+
+	public void deleteAll(Object correlationId) {
+		correlationToMessage.remove(correlationId);
+	}
+
+	private Collection<Message<?>> getMessagesInternal(Object correlationId) {
+		if (!correlationToMessage.containsKey(correlationId)) {
+			correlationToMessage.putIfAbsent(correlationId, new HashSet<Message<?>>());
+		}
+		Collection<Message<?>> collection = correlationToMessage.get(correlationId);
+		return collection;
 	}
 
 }

org.springframework.integration/src/test/java/org/springframework/integration/aggregator/CorrelatingMessageHandlerTests.java

@@ -108,8 +108,8 @@ public class CorrelatingMessageHandlerTests {
         handler.handleMessage(message2);
         storedMessages.add(message2);
 
-        verify(store).put(message1);
-        verify(store).put(message2);
+        verify(store).put(correlationKey, message1);
+        verify(store).put(correlationKey, message2);
         verify(store, times(2)).list(correlationKey);
         verify(correlationStrategy).getCorrelationKey(message1);
         verify(correlationStrategy).getCorrelationKey(message2);
@@ -164,10 +164,9 @@ public class CorrelatingMessageHandlerTests {
         assertFalse(handler.forceComplete("key"));
 
         bothMessagesHandled.await();
-        verify(store).put(message1);
-        verify(store).put(message2);
-        verify(store).delete(id1);
-        verify(store).delete(id2);
+        verify(store).put(correlationKey, message1);
+        verify(store).put(correlationKey, message2);
+        verify(store).deleteAll(correlationKey);
     }
 
     private Message<?> testMessage(String correlationKey, int sequenceNumber) {

org.springframework.integration/src/test/java/org/springframework/integration/aggregator/NewResequencerTests.java

@@ -120,6 +120,7 @@ public class NewResequencerTests {
         assertEquals(new Integer(1), reply1.getHeaders().getSequenceNumber());
         assertNotNull(reply2);
         assertEquals(new Integer(2), reply2.getHeaders().getSequenceNumber());
+        System.err.println(reply3);
         assertNull(reply3);
         // when sending the last message, the whole sequence must have been sent
         this.resequencer.handleMessage(message4);

org.springframework.integration/src/test/java/org/springframework/integration/store/SimpleMessageStoreTests.java

@@ -16,20 +16,25 @@
 
 package org.springframework.integration.store;
 
+import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThat;
+
+import java.util.Arrays;
+
 import org.junit.Test;
 import org.springframework.integration.core.Message;
 import org.springframework.integration.core.MessagingException;
 import org.springframework.integration.message.MessageBuilder;
 
-import static org.hamcrest.CoreMatchers.is;
-import static org.junit.Assert.assertThat;
-
 /**
  * @author Iwein Fuld
+ * @author Dave Syer
  */
-public class SimpleMessageStoreTest {
+public class SimpleMessageStoreTests {
 
 	@Test
+	@SuppressWarnings("unchecked")
 	public void shouldRetainMessage() {
 		SimpleMessageStore store = new SimpleMessageStore();
 		Message<String> testMessage1 = MessageBuilder.withPayload("foo").build();
@@ -54,4 +59,22 @@ public class SimpleMessageStoreTest {
 		store.put(testMessage1);
 		store.put(testMessage2);
 	}
+	
+	@Test
+	public void shouldListByCorrelation() throws Exception {
+		SimpleMessageStore store = new SimpleMessageStore();
+		Message<String> testMessage1 = MessageBuilder.withPayload("foo").build();
+		store.put("bar", testMessage1);
+		assertEquals(1, store.list("bar").size());
+	}
+
+	@Test
+	public void shouldListByCorrelationAfterAddAll() throws Exception {
+		SimpleMessageStore store = new SimpleMessageStore();
+		Message<String> testMessage1 = MessageBuilder.withPayload("foo").build();
+		Message<String> testMessage2 = MessageBuilder.withPayload("bar").build();
+		store.put("bar", Arrays.<Message<?>>asList(testMessage1, testMessage2));
+		assertEquals(2, store.list("bar").size());
+	}
+
 }