INT-1069: Add MessageGroupQueue and JDBC tests for rollback

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

@@ -12,6 +12,10 @@
  */
 package org.springframework.integration.aggregator;
 
+import java.util.Iterator;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.integration.core.Message;
@@ -21,10 +25,6 @@ import org.springframework.integration.store.MessageGroup;
 import org.springframework.integration.store.MessageGroupStore;
 import org.springframework.integration.store.SimpleMessageStore;
 
-import java.util.Iterator;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-
 /**
  * This Endpoint serves as a barrier for messages that should not be processed yet. The decision when a message can be
  * processed is delegated to a {@link org.springframework.integration.aggregator.ReleaseStrategy ReleaseStrategy}.
@@ -37,7 +37,7 @@ import java.util.concurrent.ConcurrentMap;
  *
  * @author Iwein Fuld
  */
-public class CorrelatingMessageBarrier extends AbstractMessageHandler implements MessageSource {
+public class CorrelatingMessageBarrier extends AbstractMessageHandler implements MessageSource<Object> {
 	private static final Log log = LogFactory.getLog(CorrelatingMessageBarrier.class);
 
 	private CorrelationStrategy correlationStrategy;
@@ -86,7 +86,7 @@ public class CorrelatingMessageBarrier extends AbstractMessageHandler implements
 	}
 
 
-	public Message receive() {
+	public Message<Object> receive() {
 		for (Object key : correlationLocks.keySet()) {
 			Object lock = getLock(key);
 			synchronized (lock) {
@@ -106,7 +106,9 @@ public class CorrelatingMessageBarrier extends AbstractMessageHandler implements
 						} else {
 							remove(key);
 						}
-						return nextMessage;
+						@SuppressWarnings("unchecked")
+						Message<Object> result = (Message<Object>) nextMessage;
+						return result;
 					}
 				}
 			}

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

@@ -108,7 +108,7 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 	public void setMessageStore(MessageGroupStore store) {
 		this.messageStore = store;
 		store.registerMessageGroupExpiryCallback(new MessageGroupCallback() {
-			public void execute(MessageGroup group) {
+			public void execute(MessageGroupStore messageGroupStore, MessageGroup group) {
 				forceComplete(group);
 			}
 		});

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

@@ -62,7 +62,6 @@ public abstract class AbstractMessageGroupStore implements MessageGroupStore, It
 			if (group.getTimestamp() < threshold) {
 				count++;
 				expire(group);
-				removeMessageGroup(group.getCorrelationKey());
 			}
 		}
 		return count;
@@ -76,7 +75,7 @@ public abstract class AbstractMessageGroupStore implements MessageGroupStore, It
 	
 		for (MessageGroupCallback callback : expiryCallbacks) {
 			try {
-				callback.execute(group);
+				callback.execute(this, group);
 			} catch (RuntimeException e) {
 				if (exception == null) {
 					exception = e;

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

@@ -21,6 +21,6 @@ package org.springframework.integration.store;
  */
 public interface MessageGroupCallback {
 	
-	void execute(MessageGroup group);
+	void execute(MessageGroupStore messageGroupStore, MessageGroup group);
 
 }

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

@@ -0,0 +1,192 @@
+/*
+ * Copyright 2002-2008 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;
+
+import java.util.AbstractQueue;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.TimeUnit;
+
+import org.springframework.integration.core.Message;
+
+/**
+ * A {@link BlockingQueue} that is backed by a {@link MessageGroupStore}. Can be used to ensure guaranteed delivery in
+ * the face of transaction rollback (assuming the store is transactional) and also to ensure messages are not lost if
+ * the process dies (assuming the store is durable). To use the queue across process re-starts, the same correlation key
+ * must be provided, so it needs to be unique but identifiable with a single logical instance of the queue.
+ * 
+ * @author Dave Syer
+ * @since 2.0
+ * 
+ */
+public class MessageGroupQueue extends AbstractQueue<Message<?>> implements BlockingQueue<Message<?>> {
+
+	private static final int DEFAULT_CAPACITY = Integer.MAX_VALUE;
+
+	private final MessageGroupStore messageGroupStore;
+
+	private final Object correlationKey;
+
+	private final int capacity;
+
+	// This one could be a global semaphore
+	private Object storeLock = new Object();
+
+	// This one only needs to be local
+	private Object writeLock = new Object();
+
+	// This one only needs to be local
+	private Object readLock = new Object();
+
+	public MessageGroupQueue(MessageGroupStore messageGroupStore, Object correlationKey) {
+		this(messageGroupStore, correlationKey, DEFAULT_CAPACITY);
+	}
+
+	public MessageGroupQueue(MessageGroupStore messageGroupStore, Object correlationKey, int capacity) {
+		this.messageGroupStore = messageGroupStore;
+		this.correlationKey = correlationKey;
+		this.capacity = capacity;
+	}
+
+	public Iterator<Message<?>> iterator() {
+		return getUnmarked().iterator();
+	}
+
+	public int size() {
+		return getUnmarked().size();
+	}
+
+	public boolean offer(Message<?> e) {
+		if (messageGroupStore.getMessageGroup(correlationKey).size() >= capacity) {
+			return false;
+		}
+		synchronized (storeLock) {
+			messageGroupStore.addMessageToGroup(correlationKey, e);
+		}
+		synchronized (readLock) {
+			readLock.notifyAll();
+		}
+		return true;
+	}
+
+	public Message<?> peek() {
+		Collection<Message<?>> unmarked = getUnmarked();
+		if (unmarked.isEmpty()) {
+			return null;
+		}
+		return unmarked.iterator().next();
+	}
+
+	public Message<?> poll() {
+		Message<?> result;
+		synchronized (storeLock) {
+			Collection<Message<?>> unmarked = getUnmarked();
+			if (unmarked.isEmpty()) {
+				return null;
+			}
+			result = unmarked.iterator().next();
+			messageGroupStore.removeMessageFromGroup(correlationKey, result);
+		}
+		synchronized (writeLock) {
+			writeLock.notifyAll();
+		}
+		return result;
+	}
+
+	public int drainTo(Collection<? super Message<?>> c) {
+		Collection<Message<?>> unmarked;
+		synchronized (storeLock) {
+			unmarked = getUnmarked();
+			c.addAll(unmarked);
+			messageGroupStore.markMessageGroup(messageGroupStore.getMessageGroup(correlationKey));
+		}
+		synchronized (writeLock) {
+			writeLock.notifyAll();
+		}
+		return unmarked.size();
+	}
+
+	public int drainTo(Collection<? super Message<?>> c, int maxElements) {
+		ArrayList<Message<?>> list = new ArrayList<Message<?>>();
+		synchronized (storeLock) {
+			Iterator<Message<?>> unmarked = getUnmarked().iterator();
+			for (int i = 0; i < maxElements && unmarked.hasNext(); i++) {
+				Message<?> message = unmarked.next();
+				messageGroupStore.removeMessageFromGroup(correlationKey, message);
+				list.add(message);
+			}
+		}
+		synchronized (writeLock) {
+			writeLock.notifyAll();
+		}
+		c.addAll(list);
+		return list.size();
+	}
+
+	public boolean offer(Message<?> e, long timeout, TimeUnit unit) throws InterruptedException {
+		if (!offer(e)) {
+			synchronized (writeLock) {
+				writeLock.wait(TimeUnit.MILLISECONDS.convert(timeout, unit));
+			}
+		}
+		return offer(e);
+	}
+
+	public Message<?> poll(long timeout, TimeUnit unit) throws InterruptedException {
+		Message<?> message = poll();
+		if (message != null) {
+			return message;
+		}
+		long threshold = System.currentTimeMillis() + TimeUnit.MILLISECONDS.convert(timeout, unit);
+		while (message == null && System.currentTimeMillis() < threshold) {
+			synchronized (readLock) {
+				readLock.wait(threshold - System.currentTimeMillis());
+			}
+			message = poll();
+		}
+		return message;
+	}
+
+	public void put(Message<?> e) throws InterruptedException {
+		while (!offer(e)) {
+			synchronized (writeLock) {
+				writeLock.wait();
+			}
+		}
+	}
+
+	public int remainingCapacity() {
+		return capacity - messageGroupStore.getMessageGroup(correlationKey).size();
+	}
+
+	public Message<?> take() throws InterruptedException {
+		Message<?> message = poll();
+		while (message == null) {
+			synchronized (readLock) {
+				readLock.wait();
+			}
+			message = poll();
+		}
+		return message;
+	}
+
+	private Collection<Message<?>> getUnmarked() {
+		return messageGroupStore.getMessageGroup(correlationKey).getUnmarked();
+	}
+
+}

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

@@ -54,15 +54,15 @@ public class CorrelatingMessageBarrierTest {
 
 	@Test
 	public void shouldPassMessage() {
-		Message message = testMessage();
+		Message<Object> message = testMessage();
 		barrier.handleMessage(message);
 		assertThat(barrier.receive(), is(message));
 	}
 
 	@Test
 	public void shouldRemoveKeyWithoutLockingOnEmptyQueue() throws InterruptedException {
-		Message message = testMessage();
-		Message message2 = testMessage();
+		Message<Object> message = testMessage();
+		Message<Object> message2 = testMessage();
 		barrier.handleMessage(message);
 		verify(correlationStrategy).getCorrelationKey(message);
 		assertThat(barrier.receive(), is(notNullValue()));
@@ -95,7 +95,7 @@ public class CorrelatingMessageBarrierTest {
 		}
 	}
 
-	private void sendAsynchronously(final MessageHandler handler, final Message<?> message, final CountDownLatch start, final CountDownLatch sent) {
+	private void sendAsynchronously(final MessageHandler handler, final Message<Object> message, final CountDownLatch start, final CountDownLatch sent) {
 		Executors.newSingleThreadExecutor().execute(new Runnable() {
 			public void run() {
 				try {
@@ -110,8 +110,8 @@ public class CorrelatingMessageBarrierTest {
 
 	}
 
-	private Message testMessage() {
-		return MessageBuilder.withPayload("payload").build();
+	private Message<Object> testMessage() {
+		return MessageBuilder.withPayload((Object)"payload").build();
 	}
 
 
@@ -145,6 +145,7 @@ public class CorrelatingMessageBarrierTest {
 			}
 		}
 
+		@SuppressWarnings("unused")
 		public void releaseAll() {
 			for (Semaphore semaphore : keyLocks.values()) {
 				semaphore.release();

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

@@ -22,7 +22,6 @@ import static org.junit.Assert.assertTrue;
 import org.junit.Test;
 import org.springframework.integration.core.Message;
 import org.springframework.integration.message.MessageBuilder;
-import org.springframework.integration.store.MessageGroup;
 import org.springframework.integration.store.SimpleMessageGroup;
 
 /**

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

@@ -22,7 +22,6 @@ import static org.junit.Assert.assertTrue;
 import org.junit.Test;
 import org.springframework.integration.core.Message;
 import org.springframework.integration.message.MessageBuilder;
-import org.springframework.integration.store.MessageGroup;
 import org.springframework.integration.store.SimpleMessageGroup;
 
 /**

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

@@ -0,0 +1,154 @@
+/*
+ * Copyright 2002-2008 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;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CompletionService;
+import java.util.concurrent.ExecutorCompletionService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.junit.Test;
+import org.springframework.integration.core.Message;
+import org.springframework.integration.message.StringMessage;
+
+/**
+ * @author Dave Syer
+ * @since 2.0
+ * 
+ */
+public class MessageGroupQueueTests {
+	
+	static final Log logger = LogFactory.getLog(MessageGroupQueueTests.class);
+
+	@Test
+	public void testPutAndPoll() throws Exception {
+		MessageGroupQueue queue = new MessageGroupQueue(new SimpleMessageStore(), "FOO");
+		queue.put(new StringMessage("foo"));
+		Message<?> result = queue.poll(100, TimeUnit.MILLISECONDS);
+		assertNotNull(result);
+	}
+
+	@Test
+	public void testSize() throws Exception {
+		MessageGroupQueue queue = new MessageGroupQueue(new SimpleMessageStore(), "FOO");
+		queue.put(new StringMessage("foo"));
+		assertEquals(1, queue.size());
+		queue.poll(100, TimeUnit.MILLISECONDS);
+		assertEquals(0, queue.size());
+	}
+
+	@Test
+	public void testCapacityAfterExpiry() throws Exception {
+		SimpleMessageStore messageGroupStore = new SimpleMessageStore();
+		MessageGroupQueue queue = new MessageGroupQueue(messageGroupStore, "FOO", 2);
+		queue.put(new StringMessage("foo"));
+		assertEquals(1, queue.remainingCapacity());
+		queue.put(new StringMessage("bar"));
+		assertEquals(0, queue.remainingCapacity());
+		Message<?> result = queue.poll(100, TimeUnit.MILLISECONDS);
+		assertNotNull(result);
+		assertEquals(1, queue.remainingCapacity());
+	}
+
+	@Test
+	public void testCapacityExceeded() throws Exception {
+		SimpleMessageStore messageGroupStore = new SimpleMessageStore();
+		MessageGroupQueue queue = new MessageGroupQueue(messageGroupStore, "FOO", 1);
+		queue.put(new StringMessage("foo"));
+		assertFalse(queue.offer(new StringMessage("bar"), 100, TimeUnit.MILLISECONDS));
+	}
+
+	@Test
+	public void testPutAndTake() throws Exception {
+		MessageGroupQueue queue = new MessageGroupQueue(new SimpleMessageStore(), "FOO");
+		queue.put(new StringMessage("foo"));
+		Message<?> result = queue.take();
+		assertNotNull(result);
+	}
+
+	@Test
+	public void testConcurrentAccess() throws Exception {
+
+		SimpleMessageStore messageGroupStore = new SimpleMessageStore();
+		final MessageGroupQueue queue = new MessageGroupQueue(messageGroupStore, "FOO");
+		CompletionService<Boolean> completionService = new ExecutorCompletionService<Boolean>(Executors
+				.newCachedThreadPool());
+
+		int concurrency = 30;
+		final int maxPerTask = 20;
+		final Set<String> set = new HashSet<String>();
+
+		for (int i = 0; i < concurrency; i++) {
+
+			final int big = i;		
+	
+			completionService.submit(new Callable<Boolean>() {
+				public Boolean call() throws Exception {
+					boolean result = true;
+					for (int j = 0; j < maxPerTask; j++) {
+						result &= queue.add(new StringMessage("count=" + big + ":" + j));
+						if (!result) {
+							logger.warn("Failed to add");
+						}
+					}
+					return result;
+				}
+			});
+
+			completionService.submit(new Callable<Boolean>() {
+				public Boolean call() throws Exception {
+					boolean result = true;
+					for (int j = 0; j < maxPerTask; j++) {
+						@SuppressWarnings("unchecked")
+						Message<String> item = (Message<String>) queue.poll(1, TimeUnit.SECONDS);
+						set.add(item.getPayload());
+						result &= item!=null;
+						if (!result) {
+							logger.warn("Failed to poll");
+						}
+					}
+					return result;
+				}
+			});
+			
+			messageGroupStore.expireMessageGroups(-10000);
+			
+		}
+
+		for (int j = 0; j < 2*concurrency; j++) {
+			assertTrue(completionService.take().get());
+		}
+		
+		// Ensure all items polled are unique
+		assertEquals(concurrency*maxPerTask, set.size());
+		
+		assertEquals(0, queue.size());
+		messageGroupStore.expireMessageGroups(-10000);
+		assertEquals(Integer.MAX_VALUE, queue.remainingCapacity());
+
+	}
+
+}

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

@@ -13,7 +13,7 @@
 
 package org.springframework.integration.store;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -57,8 +57,9 @@ public class MessageStoreReaperTests {
 
 		private static final List<MessageGroup> groups = new ArrayList<MessageGroup>();
 
-		public void execute(MessageGroup group) {
+		public void execute(MessageGroupStore messageGroupStore, MessageGroup group) {
 			groups.add(group);
+			messageGroupStore.removeMessageGroup(group.getCorrelationKey());
 		}
 		
 	}

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

@@ -39,7 +39,7 @@ public class MessageStoreTests {
 	public void shouldRegisterCallbacks() throws Exception {
 		TestMessageStore store = new TestMessageStore();
 		store.setExpiryCallbacks(Arrays.<MessageGroupCallback>asList(new MessageGroupCallback() {
-			public void execute(MessageGroup group) {
+			public void execute(MessageGroupStore messageGroupStore, MessageGroup group) {
 			}
 		}));
 		assertEquals(1, ((Collection<?>)ReflectionTestUtils.getField(store, "expiryCallbacks")).size());
@@ -51,8 +51,9 @@ public class MessageStoreTests {
 		TestMessageStore store = new TestMessageStore();
 		final List<String> list = new ArrayList<String>();
 		store.registerMessageGroupExpiryCallback(new MessageGroupCallback() {
-			public void execute(MessageGroup group) {
+			public void execute(MessageGroupStore messageGroupStore, MessageGroup group) {
 				list.add(group.getOne().getPayload().toString());
+				messageGroupStore.removeMessageGroup(group.getCorrelationKey());
 			}
 		});
 

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

@@ -105,7 +105,7 @@ public class SimpleMessageStoreTests {
 	public void shouldRegisterCallbacks() throws Exception {
 		SimpleMessageStore store = new SimpleMessageStore();
 		store.setExpiryCallbacks(Arrays.<MessageGroupCallback>asList(new MessageGroupCallback() {
-			public void execute(MessageGroup group) {
+			public void execute(MessageGroupStore messageGroupStore, MessageGroup group) {
 			}
 		}));
 		assertEquals(1, ((Collection<?>)ReflectionTestUtils.getField(store, "expiryCallbacks")).size());
@@ -117,8 +117,9 @@ public class SimpleMessageStoreTests {
 		SimpleMessageStore store = new SimpleMessageStore();
 		final List<String> list = new ArrayList<String>();
 		store.registerMessageGroupExpiryCallback(new MessageGroupCallback() {
-			public void execute(MessageGroup group) {
+			public void execute(MessageGroupStore messageGroupStore, MessageGroup group) {
 				list.add(group.getOne().getPayload().toString());
+				messageGroupStore.removeMessageGroup(group.getCorrelationKey());
 			}
 		});