Fixed bounding issue in SimpleMessageStore and introduced UpperBound to avoid code duplication with PriorityChannel

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

@@ -118,7 +118,7 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 	}
 
 	public CorrelatingMessageHandler(MessageGroupProcessor processor) {
-		this(new SimpleMessageStore(100), new HeaderAttributeCorrelationStrategy(MessageHeaders.CORRELATION_ID),
+		this(new SimpleMessageStore(0), new HeaderAttributeCorrelationStrategy(MessageHeaders.CORRELATION_ID),
 				new SequenceSizeCompletionStrategy(), processor);
 	}
 

org.springframework.integration/src/main/java/org/springframework/integration/channel/PriorityChannel.java

@@ -18,11 +18,10 @@ package org.springframework.integration.channel;
 
 import java.util.Comparator;
 import java.util.concurrent.PriorityBlockingQueue;
-import java.util.concurrent.Semaphore;
-import java.util.concurrent.TimeUnit;
 
 import org.springframework.integration.core.Message;
 import org.springframework.integration.core.MessagePriority;
+import org.springframework.integration.util.UpperBound;
 
 /**
  * A message channel that prioritizes messages based on a {@link Comparator}.
@@ -32,7 +31,7 @@ import org.springframework.integration.core.MessagePriority;
  */
 public class PriorityChannel extends QueueChannel {
 
-	private final Semaphore semaphore;
+	private final UpperBound upperBound;
 
 
 	/**
@@ -45,7 +44,7 @@ public class PriorityChannel extends QueueChannel {
 	public PriorityChannel(int capacity, Comparator<Message<?>> comparator) {
 		super(new PriorityBlockingQueue<Message<?>>(11,
 				(comparator != null) ? comparator : new MessagePriorityComparator()));
-		this.semaphore =  (capacity > 0) ? new Semaphore(capacity, true) : null;
+		this.upperBound = new UpperBound(capacity);
 	}
 
 	/**
@@ -77,7 +76,7 @@ public class PriorityChannel extends QueueChannel {
 
 	@Override
 	protected boolean doSend(Message<?> message, long timeout) {
-		if (!acquirePermitIfNecessary(timeout)) {
+		if (!upperBound.tryAcquire(timeout)) {
 			return false;
 		}
 		return super.doSend(message, 0);
@@ -87,32 +86,12 @@ public class PriorityChannel extends QueueChannel {
 	protected Message<?> doReceive(long timeout) {
 		Message<?> message = super.doReceive(timeout);
 		if (message != null) {
-			this.releasePermitIfNecessary();
+			upperBound.release();
 			return message;
 		}
 		return null;
 	}
 
-	private boolean acquirePermitIfNecessary(long timeoutInMilliseconds) {
-		if (this.semaphore != null) {
-			try {
-				return this.semaphore.tryAcquire(timeoutInMilliseconds, TimeUnit.MILLISECONDS);
-			}
-			catch (InterruptedException e) {
-				Thread.currentThread().interrupt();
-				return false;
-			}
-		}
-		return true;
-	}
-
-	private void releasePermitIfNecessary() {
-		if (this.semaphore != null) {
-			this.semaphore.release();
-		}
-	}
-
-
 	private static class MessagePriorityComparator implements Comparator<Message<?>> {
 
 		public int compare(Message<?> message1, Message<?> message2) {

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

@@ -17,6 +17,8 @@
 package org.springframework.integration.store;
 
 import org.springframework.integration.core.Message;
+import org.springframework.integration.core.MessagingException;
+import org.springframework.integration.util.UpperBound;
 import org.springframework.util.Assert;
 
 import java.util.ArrayList;
@@ -35,15 +37,31 @@ import java.util.concurrent.ConcurrentHashMap;
 public class SimpleMessageStore implements MessageStore {
 
     private final Map<Object, Message<?>> map;
+	private final UpperBound upperBound;
 
-
-    public SimpleMessageStore(int capacity) {
-        this.map = new ConcurrentHashMap<Object, Message<?>>(capacity);
+	/**
+	 * Creates a SimpleMessageStore with a maximum size limited by the given capacity, or unlimited
+	 * size if the given capacity is less than 1.
+	 */
+	public SimpleMessageStore(int capacity) {
+        this.map = new ConcurrentHashMap<Object, Message<?>>();
+		this.upperBound =  new UpperBound(capacity);
     }
 
+	/**
+	 * Creates a SimpleMessageStore with unlimited capacity
+	 */
+	public SimpleMessageStore() {
+		this(0);
+	}
 
-    @SuppressWarnings("unchecked")
+
+	@SuppressWarnings("unchecked")
     public <T> Message<T> put(Message<T> message) {
+		if (!upperBound.tryAcquire(0)){
+			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);
     }
 
@@ -56,7 +74,11 @@ public class SimpleMessageStore implements MessageStore {
     }
 
     public Message<?> delete(Object key) {
-        return (key != null) ? this.map.remove(key) : null;
+		if (key != null) {
+			upperBound.release();
+			return this.map.remove(key);
+		}
+		else return null;
     }
 
     public int size() {

org.springframework.integration/src/main/java/org/springframework/integration/util/UpperBound.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2002-2010 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.util;
+
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Thin wrapper around a Semaphore that allows to create a potentially unlimited upper bound
+ * to by used in buffers of messages (e.g. a QueueChannel or a MessageStore).
+ *
+ * @author Mark Fisher
+ * @author Iwein Fuld
+ * @since 2.0
+ */
+public final class UpperBound {
+	public final Semaphore semaphore;
+
+	/**
+	 * Create an UpperBound with the given capacity. If the given capacity is less than 1
+	 * an infinite UpperBound is created.
+	 *
+	 * @param capacity
+	 */
+	public UpperBound(int capacity) {
+		this.semaphore = (capacity > 0) ? new Semaphore(capacity, true) : null;
+	}
+
+	/**
+	 * Acquires a lock on the underlying semaphore if this UpperBound is bounded and returns true
+	 * if it succeeds within the given timeout
+	 */
+	public boolean tryAcquire(long timeoutInMilliseconds) {
+		if (this.semaphore != null) {
+			try {
+				return this.semaphore.tryAcquire(timeoutInMilliseconds, TimeUnit.MILLISECONDS);
+			}
+			catch (InterruptedException e) {
+				Thread.currentThread().interrupt();
+				return false;
+			}
+		}
+		return true;
+	}
+
+	/**
+	 * Releases one lock on the underlying semaphore. This is typically not done by the same Thread
+	 * that acquired the lock, but by the thread that picked up the message.
+	 */
+	public void release() {
+		if (this.semaphore != null) {
+			this.semaphore.release();
+		}
+	}
+}

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

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2002-2010 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 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
+ */
+public class SimpleMessageStoreTest {
+
+	@Test
+	public void shouldRetainMessage() {
+		SimpleMessageStore store = new SimpleMessageStore();
+		Message<String> testMessage1 = MessageBuilder.withPayload("foo").build();
+		store.put(testMessage1);
+		assertThat((Message<String>) store.get(testMessage1.getHeaders().getId()), is(testMessage1));
+	}
+
+	@Test(expected = MessagingException.class)
+	public void shouldNotHoldMoreThanCapacity() {
+		SimpleMessageStore store = new SimpleMessageStore(1);
+		Message<String> testMessage1 = MessageBuilder.withPayload("foo").build();
+		Message<String> testMessage2 = MessageBuilder.withPayload("bar").build();
+		store.put(testMessage1);
+		store.put(testMessage2);
+	}
+
+	@Test
+	public void shouldHoldCapacityExactly() {
+		SimpleMessageStore store = new SimpleMessageStore(2);
+		Message<String> testMessage1 = MessageBuilder.withPayload("foo").build();
+		Message<String> testMessage2 = MessageBuilder.withPayload("bar").build();
+		store.put(testMessage1);
+		store.put(testMessage2);
+	}
+}