GH-5123: Add LockRegistry to AbstractMessageGroupStore

Fixes: https://github.com/spring-projects/spring-integration/issues/5123

When `RedisMessageStore`, for example, adds and removes messages, it operates on two keys separately, which may cause problems in multi-threading due to non-atomic operations.
Although using Redis to delay messages is not a good idea, the abnormal loss of messages in the logs alerted me when the number of requests was not large.
By comparing the logs, the problem that the message group representing the metadata is not consistent with the actual message.

A simple solution is to add lock like in the `SimpleMessageStore`, which is also the approach taken in this pull request.

* Add `LockRegistry` to `AbstractMessageGroupStore`
* Normalize access levels and method name about the lock of `MessageGroupStore`
* Add document about the lock of `AbstractMessageGroupStore`

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2024 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -38,13 +38,12 @@ import org.springframework.util.Assert;
  * @author Gary Russell
  * @author Artem Bilan
  * @author Ngoc Nhan
+ * @author Youbin Wu
  *
  * @since 2.1
  */
 public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupStore implements MessageStore {
 
-	private static final String GROUP_ID_MUST_NOT_BE_NULL = "'groupId' must not be null";
-
 	protected static final String MESSAGE_KEY_PREFIX = "MESSAGE_";
 
 	protected static final String MESSAGE_GROUP_KEY_PREFIX = "GROUP_OF_MESSAGES_";
@@ -206,7 +205,7 @@ public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupS
 	}
 
 	@Override
-	public void addMessagesToGroup(Object groupId, Message<?>... messages) {
+	protected void doAddMessagesToGroup(Object groupId, Message<?>... messages) {
 		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
 		Assert.notNull(messages, "'messages' must not be null");
 
@@ -240,7 +239,7 @@ public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupS
 	}
 
 	@Override
-	public void removeMessagesFromGroup(Object groupId, Collection<Message<?>> messages) {
+	protected void doRemoveMessagesFromGroup(Object groupId, Collection<Message<?>> messages) {
 		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
 		Assert.notNull(messages, "'messages' must not be null");
 
@@ -283,7 +282,7 @@ public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupS
 	}
 
 	@Override
-	public boolean removeMessageFromGroupById(Object groupId, UUID messageId) {
+	protected boolean doRemoveMessageFromGroupById(Object groupId, UUID messageId) {
 		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
 		Assert.notNull(messageId, "'messageId' must not be null");
 		Object mgm = doRetrieve(this.groupPrefix + groupId);
@@ -305,7 +304,7 @@ public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupS
 	}
 
 	@Override
-	public void completeGroup(Object groupId) {
+	protected void doCompleteGroup(Object groupId) {
 		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
 		MessageGroupMetadata metadata = getGroupMetadata(groupId);
 		if (metadata != null) {
@@ -319,7 +318,7 @@ public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupS
 	 * Remove the MessageGroup with the provided group ID.
 	 */
 	@Override
-	public void removeMessageGroup(Object groupId) {
+	protected void doRemoveMessageGroup(Object groupId) {
 		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
 		Object mgm = doRemove(this.groupPrefix + groupId);
 		if (mgm != null) {
@@ -337,7 +336,7 @@ public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupS
 	}
 
 	@Override
-	public void setGroupCondition(Object groupId, String condition) {
+	protected void doSetGroupCondition(Object groupId, String condition) {
 		MessageGroupMetadata metadata = getGroupMetadata(groupId);
 		if (metadata != null) {
 			metadata.setCondition(condition);
@@ -346,7 +345,7 @@ public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupS
 	}
 
 	@Override
-	public void setLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber) {
+	protected void doSetLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber) {
 		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
 		MessageGroupMetadata metadata = getGroupMetadata(groupId);
 		if (metadata == null) {
@@ -359,7 +358,7 @@ public abstract class AbstractKeyValueMessageStore extends AbstractMessageGroupS
 	}
 
 	@Override
-	public Message<?> pollMessageFromGroup(Object groupId) {
+	protected Message<?> doPollMessageFromGroup(Object groupId) {
 		MessageGroupMetadata groupMetadata = getGroupMetadata(groupId);
 		if (groupMetadata != null) {
 			UUID firstId = groupMetadata.firstId();

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -19,16 +19,22 @@ package org.springframework.integration.store;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.LinkedHashSet;
+import java.util.UUID;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+import org.springframework.integration.support.locks.DefaultLockRegistry;
+import org.springframework.integration.support.locks.LockRegistry;
+import org.springframework.integration.util.CheckedCallable;
+import org.springframework.integration.util.CheckedRunnable;
 import org.springframework.jmx.export.annotation.ManagedAttribute;
 import org.springframework.jmx.export.annotation.ManagedOperation;
 import org.springframework.jmx.export.annotation.ManagedResource;
 import org.springframework.messaging.Message;
+import org.springframework.util.Assert;
 
 /**
  * @author Dave Syer
@@ -36,6 +42,7 @@ import org.springframework.messaging.Message;
  * @author Gary Russell
  * @author Artem Bilan
  * @author Christian Tzolov
+ * @author Youbin Wu
  *
  * @since 2.0
  */
@@ -43,6 +50,10 @@ import org.springframework.messaging.Message;
 public abstract class AbstractMessageGroupStore extends AbstractBatchingMessageGroupStore
 		implements MessageGroupStore, Iterable<MessageGroup> {
 
+	protected static final String INTERRUPTED_WHILE_OBTAINING_LOCK = "Interrupted while obtaining lock";
+
+	protected static final String GROUP_ID_MUST_NOT_BE_NULL = "'groupId' must not be null";
+
 	protected final Log logger = LogFactory.getLog(getClass()); // NOSONAR final
 
 	private final Lock lock = new ReentrantLock();
@@ -56,6 +67,8 @@ public abstract class AbstractMessageGroupStore extends AbstractBatchingMessageG
 
 	private boolean timeoutOnIdle;
 
+	private LockRegistry lockRegistry = new DefaultLockRegistry();
+
 	protected AbstractMessageGroupStore() {
 	}
 
@@ -109,6 +122,20 @@ public abstract class AbstractMessageGroupStore extends AbstractBatchingMessageG
 		this.lazyLoadMessageGroups = lazyLoadMessageGroups;
 	}
 
+	/**
+	 * Specify the type of the {@link LockRegistry} to ensure atomic operations
+	 * @param lockRegistry lockRegistryType
+	 * @since 6.5
+	 */
+	public final void setLockRegistry(LockRegistry lockRegistry) {
+		Assert.notNull(lockRegistry, "The LockRegistry cannot be null");
+		this.lockRegistry = lockRegistry;
+	}
+
+	protected LockRegistry getLockRegistry() {
+		return this.lockRegistry;
+	}
+
 	@Override
 	public void registerMessageGroupExpiryCallback(MessageGroupCallback callback) {
 		if (callback instanceof UniqueExpiryCallback) {
@@ -195,12 +222,98 @@ public abstract class AbstractMessageGroupStore extends AbstractBatchingMessageG
 		removeMessagesFromGroup(key, Arrays.asList(messages));
 	}
 
+	@Override
+	public void removeMessagesFromGroup(Object key, Collection<Message<?>> messages) {
+		Assert.notNull(key, GROUP_ID_MUST_NOT_BE_NULL);
+		executeLocked(key, () -> doRemoveMessagesFromGroup(key, messages));
+	}
+
+	protected abstract void doRemoveMessagesFromGroup(Object key, Collection<Message<?>> messages);
+
+	@Override
+	public void addMessagesToGroup(Object groupId, Message<?>... messages) {
+		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
+		executeLocked(groupId, () -> doAddMessagesToGroup(groupId, messages));
+	}
+
+	protected abstract void doAddMessagesToGroup(Object groupId, Message<?>... messages);
+
 	@Override
 	public MessageGroup addMessageToGroup(Object groupId, Message<?> message) {
 		addMessagesToGroup(groupId, message);
 		return getMessageGroup(groupId);
 	}
 
+	@Override
+	public void removeMessageGroup(Object groupId) {
+		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
+		executeLocked(groupId, () -> doRemoveMessageGroup(groupId));
+	}
+
+	protected abstract void doRemoveMessageGroup(Object groupId);
+
+	@Override
+	public boolean removeMessageFromGroupById(Object groupId, UUID messageId) {
+		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
+		return executeLocked(groupId, () -> doRemoveMessageFromGroupById(groupId, messageId));
+	}
+
+	protected boolean doRemoveMessageFromGroupById(Object groupId, UUID messageId) {
+		throw new UnsupportedOperationException("Not supported for this store");
+	}
+
+	@Override
+	public void setLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber) {
+		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
+		executeLocked(groupId, () -> doSetLastReleasedSequenceNumberForGroup(groupId, sequenceNumber));
+	}
+
+	protected abstract void doSetLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber);
+
+	@Override
+	public void completeGroup(Object groupId) {
+		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
+		executeLocked(groupId, () -> doCompleteGroup(groupId));
+	}
+
+	protected abstract void doCompleteGroup(Object groupId);
+
+	@Override
+	public void setGroupCondition(Object groupId, String condition) {
+		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
+		executeLocked(groupId, () -> doSetGroupCondition(groupId, condition));
+	}
+
+	protected abstract void doSetGroupCondition(Object groupId, String condition);
+
+	@Override
+	public Message<?> pollMessageFromGroup(Object groupId) {
+		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
+		return executeLocked(groupId, () -> doPollMessageFromGroup(groupId));
+	}
+
+	protected abstract Message<?> doPollMessageFromGroup(Object groupId);
+
+	protected <T, E extends RuntimeException> T executeLocked(Object groupId, CheckedCallable<T, E> runnable) {
+		try {
+			return this.lockRegistry.executeLocked(groupId, runnable);
+		}
+		catch (InterruptedException ex) {
+			Thread.currentThread().interrupt();
+			throw new IllegalStateException(INTERRUPTED_WHILE_OBTAINING_LOCK, ex);
+		}
+	}
+
+	protected <E extends RuntimeException> void executeLocked(Object groupId, CheckedRunnable<E> runnable) {
+		try {
+			this.lockRegistry.executeLocked(groupId, runnable);
+		}
+		catch (InterruptedException ex) {
+			Thread.currentThread().interrupt();
+			throw new IllegalStateException(INTERRUPTED_WHILE_OBTAINING_LOCK, ex);
+		}
+	}
+
 	private void expire(MessageGroup group) {
 
 		RuntimeException exception = null;

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2024 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -45,6 +45,7 @@ import org.springframework.util.CollectionUtils;
  * @author Gary Russell
  * @author Ryan Barker
  * @author Artem Bilan
+ * @author Youbin Wu
  *
  * @since 2.0
  */
@@ -55,8 +56,6 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 
 	private static final String UPPER_BOUND_MUST_NOT_BE_NULL = "'upperBound' must not be null.";
 
-	private static final String INTERRUPTED_WHILE_OBTAINING_LOCK = "Interrupted while obtaining lock";
-
 	private final ConcurrentMap<UUID, Message<?>> idToMessage = new ConcurrentHashMap<>();
 
 	private final ConcurrentMap<Object, MessageGroup> groupIdToMessageGroup = new ConcurrentHashMap<>();
@@ -71,12 +70,8 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 
 	private final long upperBoundTimeout;
 
-	private LockRegistry lockRegistry;
-
 	private boolean copyOnGet = false;
 
-	private volatile boolean isUsed;
-
 	/**
 	 * Creates a SimpleMessageStore with a maximum size limited by the given capacity, or unlimited size if the given
 	 * capacity is less than 1. The capacities are applied independently to messages stored via
@@ -133,8 +128,8 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 		this.individualUpperBound = new UpperBound(individualCapacity);
 		this.individualCapacity = individualCapacity;
 		this.groupCapacity = groupCapacity;
-		this.lockRegistry = lockRegistry;
 		this.upperBoundTimeout = upperBoundTimeout;
+		setLockRegistry(lockRegistry);
 	}
 
 	/**
@@ -162,12 +157,6 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 		this.copyOnGet = copyOnGet;
 	}
 
-	public void setLockRegistry(LockRegistry lockRegistry) {
-		Assert.notNull(lockRegistry, "The LockRegistry cannot be null");
-		Assert.isTrue(!(this.isUsed), "Cannot change the lock registry after the store has been used");
-		this.lockRegistry = lockRegistry;
-	}
-
 	@Override
 	public void setLazyLoadMessageGroups(boolean lazyLoadMessageGroups) {
 		throw new UnsupportedOperationException("The lazy-load isn't supported for in-memory 'SimpleMessageStore'");
@@ -181,7 +170,6 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 
 	@Override
 	public <T> Message<T> addMessage(Message<T> message) {
-		this.isUsed = true;
 		if (!this.individualUpperBound.tryAcquire(this.upperBoundTimeout)) {
 			throw new MessagingException(getClass().getSimpleName()
 					+ " was out of capacity ("
@@ -246,7 +234,7 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 	@Override
 	protected MessageGroup copy(MessageGroup group) {
 		Object groupId = group.getGroupId();
-		Lock lock = this.lockRegistry.obtain(groupId);
+		Lock lock = getLockRegistry().obtain(groupId);
 		try {
 			lock.lockInterruptibly();
 			try {
@@ -261,9 +249,9 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 				lock.unlock();
 			}
 		}
-		catch (InterruptedException e) {
+		catch (InterruptedException ex) {
 			Thread.currentThread().interrupt();
-			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, e);
+			throw new IllegalStateException(INTERRUPTED_WHILE_OBTAINING_LOCK, ex);
 		}
 	}
 
@@ -272,7 +260,7 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 		Assert.notNull(groupId, "'groupId' must not be null");
 		Assert.notNull(messages, "'messages' must not be null");
 
-		Lock lock = this.lockRegistry.obtain(groupId);
+		Lock lock = getLockRegistry().obtain(groupId);
 		try {
 			lock.lockInterruptibly();
 			boolean unlocked = false;
@@ -314,71 +302,50 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 				}
 			}
 		}
-		catch (InterruptedException e) {
+		catch (InterruptedException ex) {
 			Thread.currentThread().interrupt();
-			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, e);
+			throw new IllegalStateException(INTERRUPTED_WHILE_OBTAINING_LOCK, ex);
 		}
 	}
 
-	private MessagingException outOfCapacityException(Object groupId) {
-		return new MessagingException(getClass().getSimpleName() +
+	@Override
+	protected void doAddMessagesToGroup(Object groupId, Message<?>... messages) {
+		// No implementation: the addMessagesToGroup() fully uses locking algorithm.
+	}
+
+	private IllegalStateException outOfCapacityException(Object groupId) {
+		return new IllegalStateException(getClass().getSimpleName() +
 				" was out of capacity (" + this.groupCapacity + ") for group '" + groupId +
 				"', try constructing it with a larger number.");
 	}
 
 	@Override
-	public void removeMessageGroup(Object groupId) {
-		Lock lock = this.lockRegistry.obtain(groupId);
-		try {
-			lock.lockInterruptibly();
-			try {
-				MessageGroup messageGroup = this.groupIdToMessageGroup.remove(groupId);
-				if (messageGroup != null) {
-					UpperBound upperBound = this.groupToUpperBound.remove(groupId);
-					Assert.state(upperBound != null, UPPER_BOUND_MUST_NOT_BE_NULL);
-					upperBound.release(this.groupCapacity);
-				}
-			}
-			finally {
-				lock.unlock();
-			}
-		}
-		catch (InterruptedException e) {
-			Thread.currentThread().interrupt();
-			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, e);
+	protected void doRemoveMessageGroup(Object groupId) {
+		MessageGroup messageGroup = this.groupIdToMessageGroup.remove(groupId);
+		if (messageGroup != null) {
+			UpperBound upperBound = this.groupToUpperBound.remove(groupId);
+			Assert.state(upperBound != null, UPPER_BOUND_MUST_NOT_BE_NULL);
+			upperBound.release(this.groupCapacity);
 		}
 	}
 
 	@Override
-	public void removeMessagesFromGroup(Object groupId, Collection<Message<?>> messages) {
-		Lock lock = this.lockRegistry.obtain(groupId);
-		try {
-			lock.lockInterruptibly();
-			try {
-				MessageGroup group = this.groupIdToMessageGroup.get(groupId);
-				Assert.notNull(group,
-						() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' " +
-								"can not be located while attempting to remove Message(s) from the MessageGroup");
-				UpperBound upperBound = this.groupToUpperBound.get(groupId);
-				Assert.state(upperBound != null, UPPER_BOUND_MUST_NOT_BE_NULL);
-				boolean modified = false;
-				for (Message<?> messageToRemove : messages) {
-					if (group.remove(messageToRemove)) {
-						upperBound.release();
-						modified = true;
-					}
-				}
-				if (modified) {
-					group.setLastModified(System.currentTimeMillis());
-				}
-			}
-			finally {
-				lock.unlock();
+	protected void doRemoveMessagesFromGroup(Object groupId, Collection<Message<?>> messages) {
+		MessageGroup group = this.groupIdToMessageGroup.get(groupId);
+		Assert.notNull(group,
+				() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' " +
+						"can not be located while attempting to remove Message(s) from the MessageGroup");
+		UpperBound upperBound = this.groupToUpperBound.get(groupId);
+		Assert.state(upperBound != null, UPPER_BOUND_MUST_NOT_BE_NULL);
+		boolean modified = false;
+		for (Message<?> messageToRemove : messages) {
+			if (group.remove(messageToRemove)) {
+				upperBound.release();
+				modified = true;
 			}
 		}
-		catch (InterruptedException e) {
-			Thread.currentThread().interrupt();
-			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, e);
+		if (modified) {
+			group.setLastModified(System.currentTimeMillis());
 		}
 	}
 
@@ -397,35 +364,22 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 	}
 
 	@Override
-	public boolean removeMessageFromGroupById(Object groupId, UUID messageId) {
-		Lock lock = this.lockRegistry.obtain(groupId);
-		try {
-			lock.lockInterruptibly();
-			try {
-				MessageGroup group = this.groupIdToMessageGroup.get(groupId);
-				Assert.notNull(group,
-						() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' " +
-								"can not be located while attempting to remove Message from the MessageGroup");
-				UpperBound upperBound = this.groupToUpperBound.get(groupId);
-				Assert.state(upperBound != null, UPPER_BOUND_MUST_NOT_BE_NULL);
-				for (Message<?> message : group.getMessages()) {
-					if (messageId.equals(message.getHeaders().getId())) {
-						group.remove(message);
-						upperBound.release();
-						group.setLastModified(System.currentTimeMillis());
-						return true;
-					}
-				}
-				return false;
-			}
-			finally {
-				lock.unlock();
+	protected boolean doRemoveMessageFromGroupById(Object groupId, UUID messageId) {
+		MessageGroup group = this.groupIdToMessageGroup.get(groupId);
+		Assert.notNull(group,
+				() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' " +
+						"can not be located while attempting to remove Message from the MessageGroup");
+		UpperBound upperBound = this.groupToUpperBound.get(groupId);
+		Assert.state(upperBound != null, UPPER_BOUND_MUST_NOT_BE_NULL);
+		for (Message<?> message : group.getMessages()) {
+			if (messageId.equals(message.getHeaders().getId())) {
+				group.remove(message);
+				upperBound.release();
+				group.setLastModified(System.currentTimeMillis());
+				return true;
 			}
 		}
-		catch (InterruptedException ex) {
-			Thread.currentThread().interrupt();
-			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, ex);
-		}
+		return false;
 	}
 
 	@Override
@@ -434,7 +388,7 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 	}
 
 	@Override
-	public void setGroupCondition(Object groupId, String condition) {
+	protected void doSetGroupCondition(Object groupId, String condition) {
 		MessageGroup group = this.groupIdToMessageGroup.get(groupId);
 		if (group != null) {
 			group.setCondition(condition);
@@ -442,53 +396,27 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 	}
 
 	@Override
-	public void setLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber) {
-		Lock lock = this.lockRegistry.obtain(groupId);
-		try {
-			lock.lockInterruptibly();
-			try {
-				MessageGroup group = this.groupIdToMessageGroup.get(groupId);
-				Assert.notNull(group,
-						() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' " +
-								"can not be located while attempting to set 'lastReleasedSequenceNumber'");
-				group.setLastReleasedMessageSequenceNumber(sequenceNumber);
-				group.setLastModified(System.currentTimeMillis());
-			}
-			finally {
-				lock.unlock();
-			}
-		}
-		catch (InterruptedException e) {
-			Thread.currentThread().interrupt();
-			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, e);
-		}
+	protected void doSetLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber) {
+		MessageGroup group = this.groupIdToMessageGroup.get(groupId);
+		Assert.notNull(group,
+				() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' " +
+						"can not be located while attempting to set 'lastReleasedSequenceNumber'");
+		group.setLastReleasedMessageSequenceNumber(sequenceNumber);
+		group.setLastModified(System.currentTimeMillis());
 	}
 
 	@Override
-	public void completeGroup(Object groupId) {
-		Lock lock = this.lockRegistry.obtain(groupId);
-		try {
-			lock.lockInterruptibly();
-			try {
-				MessageGroup group = this.groupIdToMessageGroup.get(groupId);
-				Assert.notNull(group,
-						() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' " +
-								"can not be located while attempting to complete the MessageGroup");
-				group.complete();
-				group.setLastModified(System.currentTimeMillis());
-			}
-			finally {
-				lock.unlock();
-			}
-		}
-		catch (InterruptedException e) {
-			Thread.currentThread().interrupt();
-			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, e);
-		}
+	protected void doCompleteGroup(Object groupId) {
+		MessageGroup group = this.groupIdToMessageGroup.get(groupId);
+		Assert.notNull(group,
+				() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' " +
+						"can not be located while attempting to complete the MessageGroup");
+		group.complete();
+		group.setLastModified(System.currentTimeMillis());
 	}
 
 	@Override
-	public Message<?> pollMessageFromGroup(Object groupId) {
+	protected Message<?> doPollMessageFromGroup(Object groupId) {
 		Collection<Message<?>> messageList = getMessageGroup(groupId).getMessages();
 		Message<?> message = null;
 		if (!CollectionUtils.isEmpty(messageList)) {
@@ -521,7 +449,7 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 	}
 
 	public void clearMessageGroup(Object groupId) {
-		Lock lock = this.lockRegistry.obtain(groupId);
+		Lock lock = getLockRegistry().obtain(groupId);
 		try {
 			lock.lockInterruptibly();
 			try {
@@ -539,9 +467,9 @@ public class SimpleMessageStore extends AbstractMessageGroupStore
 				lock.unlock();
 			}
 		}
-		catch (InterruptedException e) {
+		catch (InterruptedException ex) {
 			Thread.currentThread().interrupt();
-			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, e);
+			throw new IllegalStateException(INTERRUPTED_WHILE_OBTAINING_LOCK, ex);
 		}
 	}
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -56,6 +56,7 @@ import static org.mockito.Mockito.verify;
  * @author Gary Russell
  * @author Artem Bilan
  * @author Meherzad Lahewala
+ * @author Youbin Wu
  *
  * @since 2.2
  *
@@ -350,7 +351,7 @@ public class AbstractCorrelatingMessageHandlerTests {
 		SimpleMessageStore messageStore = new SimpleMessageStore() {
 
 			@Override
-			public void removeMessageGroup(Object groupId) {
+			protected void doRemoveMessageGroup(Object groupId) {
 				throw new RuntimeException("intentional");
 			}
 		};

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2024 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -22,9 +22,8 @@ import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 
-import org.junit.Test;
+import org.junit.jupiter.api.Test;
 
-import org.springframework.integration.store.MessageGroupStore.MessageGroupCallback;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.support.GenericMessage;
 import org.springframework.test.util.ReflectionTestUtils;
@@ -35,22 +34,22 @@ import static org.assertj.core.api.Assertions.assertThat;
  * @author Dave Syer
  * @author Gary Russell
  * @author Artem Bilan
+ * @author Youbin Wu
  */
 public class MessageStoreTests {
 
 	@Test
-	public void shouldRegisterCallbacks() throws Exception {
+	public void shouldRegisterCallbacks() {
 		TestMessageStore store = new TestMessageStore();
-		store.setExpiryCallbacks(Collections.<MessageGroupCallback>singletonList((messageGroupStore, group) -> {
+		store.setExpiryCallbacks(Collections.singletonList((messageGroupStore, group) -> {
 		}));
 		assertThat(((Collection<?>) ReflectionTestUtils.getField(store, "expiryCallbacks")).size()).isEqualTo(1);
 	}
 
 	@Test
-	public void shouldExpireMessageGroup() throws Exception {
-
+	public void shouldExpireMessageGroup() {
 		TestMessageStore store = new TestMessageStore();
-		final List<String> list = new ArrayList<String>();
+		final List<String> list = new ArrayList<>();
 		store.registerMessageGroupExpiryCallback((messageGroupStore, group) -> {
 			list.add(group.getOne().getPayload().toString());
 			messageGroupStore.removeMessageGroup(group.getGroupId());
@@ -63,13 +62,13 @@ public class MessageStoreTests {
 	}
 
 	@Test
-	public void testGroupCount() throws Exception {
+	public void testGroupCount() {
 		TestMessageStore store = new TestMessageStore();
 		assertThat(store.getMessageGroupCount()).isEqualTo(1);
 	}
 
 	@Test
-	public void testGroupSizes() throws Exception {
+	public void testGroupSizes() {
 		TestMessageStore store = new TestMessageStore();
 		assertThat(store.getMessageCountForAllMessageGroups()).isEqualTo(1);
 	}
@@ -91,7 +90,7 @@ public class MessageStoreTests {
 		}
 
 		@Override
-		public void addMessagesToGroup(Object groupId, Message<?>... messages) {
+		protected void doAddMessagesToGroup(Object groupId, Message<?>... messages) {
 			throw new UnsupportedOperationException();
 		}
 
@@ -101,30 +100,30 @@ public class MessageStoreTests {
 		}
 
 		@Override
-		public void removeMessagesFromGroup(Object key, Collection<Message<?>> messages) {
+		protected void doRemoveMessagesFromGroup(Object key, Collection<Message<?>> messages) {
 			throw new UnsupportedOperationException();
 		}
 
 		@Override
-		public void removeMessageGroup(Object correlationKey) {
+		protected void doRemoveMessageGroup(Object correlationKey) {
 			if (correlationKey.equals(testMessages.getGroupId())) {
 				removed = true;
 			}
 		}
 
 		@Override
-		public void setLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber) {
+		protected void doSetLastReleasedSequenceNumberForGroup(Object groupId, int sequenceNumber) {
 			throw new UnsupportedOperationException();
 		}
 
 		@Override
-		public void completeGroup(Object groupId) {
+		protected void doCompleteGroup(Object groupId) {
 
 			throw new UnsupportedOperationException();
 		}
 
 		@Override
-		public Message<?> pollMessageFromGroup(Object groupId) {
+		protected Message<?> doPollMessageFromGroup(Object groupId) {
 			return null;
 		}
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2024 the original author or authors.
+ * Copyright 2002-2025 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.
@@ -34,6 +34,7 @@ import org.springframework.test.util.ReflectionTestUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
 
 /**
  * @author Iwein Fuld
@@ -45,7 +46,6 @@ import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 public class SimpleMessageStoreTests {
 
 	@Test
-	@SuppressWarnings("unchecked")
 	public void shouldRetainMessage() {
 		SimpleMessageStore store = new SimpleMessageStore();
 		Message<String> testMessage1 = MessageBuilder.withPayload("foo").build();
@@ -124,7 +124,7 @@ public class SimpleMessageStoreTests {
 		Message<String> testMessage1 = MessageBuilder.withPayload("foo").build();
 		Message<String> testMessage2 = MessageBuilder.withPayload("bar").build();
 		store.addMessageToGroup("foo", testMessage1);
-		assertThatExceptionOfType(MessagingException.class)
+		assertThatIllegalStateException()
 				.isThrownBy(() -> store.addMessageToGroup("foo", testMessage2));
 	}
 
@@ -158,7 +158,7 @@ public class SimpleMessageStoreTests {
 		SimpleMessageStore store2 = new SimpleMessageStore(1, 1, 1);
 		store2.addMessageToGroup("foo", MessageBuilder.withPayload("foo").build());
 
-		assertThatExceptionOfType(MessagingException.class)
+		assertThatIllegalStateException()
 				.isThrownBy(() -> store2.addMessageToGroup("foo", MessageBuilder.withPayload("bar").build()));
 	}
 
@@ -178,13 +178,13 @@ public class SimpleMessageStoreTests {
 		Message<String> testMessage2 = MessageBuilder.withPayload("bar").build();
 		store.addMessageToGroup("foo", testMessage1);
 
-		assertThatExceptionOfType(MessagingException.class)
+		assertThatIllegalStateException()
 				.isThrownBy(() -> store.addMessageToGroup("foo", testMessage2))
 				.withMessageContaining("was out of capacity (1) for group 'foo'");
 
 		store.removeMessagesFromGroup("foo", testMessage2);
 
-		assertThatExceptionOfType(MessagingException.class)
+		assertThatIllegalStateException()
 				.isThrownBy(() -> store.addMessageToGroup("foo", testMessage2))
 				.withMessageContaining("was out of capacity (1) for group 'foo'");