INT-1174: prepared MessageGroup to support marking single message

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

@@ -55,7 +55,7 @@ public interface MessageGroup {
 	 * Mark all unmarked messages in the group. A MessageGroupProcessor typically invokes this method after
 	 * processing all unmarked messages.
 	 */
-	void mark();
+	void markAll();
 
 	/**
 	 * @return a single message from the group

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

@@ -44,21 +44,29 @@ public interface MessageGroupStore {
 	/**
 	 * Persist the mark on all the messages from the group. The group is modified in the process as all its unmarked
 	 * messages become marked.
-	 * 
+	 *
 	 * @param group a MessageGroup with no unmarked messages
 	 */
 	void markMessageGroup(MessageGroup group);
 
+	/**
+	 * Persist a mark on a single message from the group. The group is modified to reflect that 'messageToMark' is no
+	 * longer unmarked but became marked instead.
+	 * @param key the correlationKey for the group containing the message
+	 * @param messageToMark the message to be marked
+	 */
+	void markMessageInGroup(Object key, Message<?> messageToMark);
+
 	/**
 	 * Remove the message group with this correlation key.
-	 * 
+	 *
 	 * @param correlationKey the correlation id to remove
 	 */
 	void removeMessageGroup(Object correlationKey);
 
 	/**
 	 * Register a callback for when a message group is expired through {@link #expireMessageGroups(long)}.
-	 * 
+	 *
 	 * @param callback a callback to execute when a message group is cleaned up
 	 */
 	void registerMessageGroupExpiryCallback(MessageGroupCallback callback);
@@ -68,12 +76,11 @@ public interface MessageGroupStore {
 	 * each of the registered callbacks on them in turn. For example: call with a timeout of 100 to expire all groups
 	 * that were created more than 100 milliseconds ago, and are not yet complete. Use a timeout of 0 (or negative to be
 	 * on the safe side) to expire all message groups.
-	 * 
+	 *
 	 * @param timeout the timeout threshold to use
 	 * @return the number of message groups expired
-	 * 
+	 *
 	 * @see #registerMessageGroupExpiryCallback(MessageGroupCallback)
 	 */
 	int expireMessageGroups(long timeout);
-
 }

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

@@ -129,7 +129,14 @@ public class SimpleMessageGroup implements MessageGroup {
 		return getOne().getHeaders().getSequenceSize();
 	}
 
-	public void mark() {
+	public void mark(Message<?> messageToMark) {
+		synchronized (lock) {
+			unmarked.remove(messageToMark);
+			marked.offer(messageToMark);
+		}
+	}
+
+	public void markAll() {
 		synchronized (lock) {
 			unmarked.drainTo(marked);
 		}
@@ -184,5 +191,4 @@ public class SimpleMessageGroup implements MessageGroup {
 		}
 		return false;
 	}
-
 }

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

@@ -113,8 +113,8 @@ public class SimpleMessageStore extends AbstractMessageGroupStore implements Mes
 	public void markMessageGroup(MessageGroup group) {
 		Object correlationId = group.getCorrelationKey();
 		MessageGroup internal = getMessageGroupInternal(correlationId);
-		internal.mark();
-		group.mark();
+		internal.markAll();
+		group.markAll();
 	}
 
 	public void removeMessageGroup(Object correlationId) {
@@ -122,6 +122,10 @@ public class SimpleMessageStore extends AbstractMessageGroupStore implements Mes
 		correlationToMessageGroup.remove(correlationId);
 	}
 
+	public void markMessageInGroup(Object key, Message<?> messageToMark) {
+		getMessageGroupInternal(key).mark(messageToMark);
+	}
+
 	@Override
 	public Iterator<MessageGroup> iterator() {
 		return new HashSet<MessageGroup>(correlationToMessageGroup.values()).iterator();

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

@@ -16,19 +16,15 @@
 
 package org.springframework.integration.store;
 
-import static org.junit.Assert.assertEquals;
-
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Iterator;
-import java.util.List;
-
 import org.junit.Test;
 import org.springframework.integration.core.Message;
 import org.springframework.integration.message.StringMessage;
 import org.springframework.test.util.ReflectionTestUtils;
 
+import java.util.*;
+
+import static org.junit.Assert.assertEquals;
+
 /**
  * @author Dave Syer
  */
@@ -82,6 +78,9 @@ public class MessageStoreTests {
 		public void markMessageGroup(MessageGroup group) {
 		}
 
+		public void markMessageInGroup(Object key, Message<?> messageToMark) {
+		}
+
 		public void removeMessageGroup(Object correlationKey) {
 			if (correlationKey.equals(testMessages.getCorrelationKey())) {
 				removed = true;

spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/JdbcMessageStore.java

@@ -76,6 +76,8 @@ public class JdbcMessageStore extends AbstractMessageGroupStore implements Messa
 
 	private static final String MARK_MESSAGES_IN_GROUP = "UPDATE %PREFIX%MESSAGE_GROUP set UPDATED_DATE=?, MARKED=1 where MARKED=0 and CORRELATION_KEY=? and REGION=?";
 
+	private static final String MARK_MESSAGE_IN_GROUP = "UPDATE %PREFIX%MESSAGE_GROUP set UPDATED_DATE=?, MARKED=1 where MARKED=0 and CORRELATION_KEY=? and REGION=? and MESSAGE_ID=?";
+
 	private static final String DELETE_MESSAGE_GROUP = "DELETE from %PREFIX%MESSAGE_GROUP where CORRELATION_KEY=? and REGION=?";
 
 	private static final String CREATE_MESSAGE_IN_GROUP = "INSERT into %PREFIX%MESSAGE_GROUP(MESSAGE_ID, REGION, CREATED_DATE, CORRELATION_KEY, MARKED, MESSAGE_BYTES)"
@@ -290,10 +292,26 @@ public class JdbcMessageStore extends AbstractMessageGroupStore implements Messa
 			}
 		});
 
-		group.mark();
+		group.markAll();
 
 	}
 
+	public void markMessageInGroup(Object key, Message<?> messageToMark) {
+		final long updatedDate = System.currentTimeMillis();
+		final String correlationId = getKey(key);
+		final String messageId = getKey(messageToMark.getHeaders().getId());
+
+		jdbcTemplate.update(getQuery(MARK_MESSAGE_IN_GROUP), new PreparedStatementSetter() {
+			public void setValues(PreparedStatement ps) throws SQLException {
+				logger.debug("Marking messages with correlation key=" + correlationId);
+				ps.setTimestamp(1, new Timestamp(updatedDate));
+				ps.setString(2, correlationId);
+				ps.setString(3, region);
+				ps.setString(4, messageId);
+			}
+		});
+	}
+
 	public void removeMessageGroup(Object correlationKey) {
 
 		final String correlationId = getKey(correlationKey);