GH-8773: Fix MGS for removal from group

Fixes #8773

The #8732 introduced a filtering for messages in group.
So, plain `removeMessage()` doesn't work any more if message is connected to some group yet.
Therefore, `DelayHandler` is failing.

* Introduce `getMessageFromGroup()` and `removeMessageFromGroupById()` into `MessageGroupStore` API
and implement it respectively in all the stores
* Remove `@LongRunningTest` from delayer integration tests and adjust its config to delay not for a long

**Cherry-pick to `6.1.x`**

Author: artembilan

spring-integration-core/src/main/java/org/springframework/integration/handler/DelayHandler.java

@@ -19,7 +19,6 @@
 import java.io.Serial;
 import java.io.Serializable;
 import java.time.Instant;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.Date;
 import java.util.List;
@@ -45,7 +44,6 @@
 import org.springframework.integration.expression.ExpressionUtils;
 import org.springframework.integration.store.MessageGroup;
 import org.springframework.integration.store.MessageGroupStore;
-import org.springframework.integration.store.MessageStore;
 import org.springframework.integration.store.SimpleMessageStore;
 import org.springframework.integration.support.management.IntegrationManagedResource;
 import org.springframework.jmx.export.annotation.ManagedResource;
@@ -115,8 +113,6 @@ public class DelayHandler extends AbstractReplyProducingMessageHandler implement
 
 	private final ConcurrentMap<String, AtomicInteger> deliveries = new ConcurrentHashMap<>();
 
-	private final Lock removeReleasedMessageLock = new ReentrantLock();
-
 	private String messageGroupId;
 
 	private long defaultDelay;
@@ -337,9 +333,7 @@ protected void doInit() {
 		if (this.messageStore == null) {
 			this.messageStore = new SimpleMessageStore();
 		}
-		else {
-			Assert.isInstanceOf(MessageStore.class, this.messageStore);
-		}
+
 		this.evaluationContext = ExpressionUtils.createStandardEvaluationContext(getBeanFactory());
 		this.releaseHandler = createReleaseMessageTask();
 	}
@@ -501,7 +495,7 @@ private Runnable releaseTaskForMessage(Message<?> delayedMessage) {
 	}
 
 	private Message<?> getMessageById(UUID messageId) {
-		Message<?> theMessage = ((MessageStore) this.messageStore).getMessage(messageId);
+		Message<?> theMessage = this.messageStore.getMessageFromGroup(this.messageGroupId, messageId);
 
 		if (theMessage == null) {
 			logger.debug(() -> "No message in the Message Store for id: " + messageId +
@@ -574,11 +568,9 @@ protected void rescheduleAt(Message<?> message, Date startTime) {
 	}
 
 	private void doReleaseMessage(Message<?> message) {
-		if (removeDelayedMessageFromMessageStore(message)
+		if (this.messageStore.removeMessageFromGroupById(this.messageGroupId, message.getHeaders().getId())
 				|| this.deliveries.get(ObjectUtils.getIdentityHexString(message)).get() > 0) {
-			if (!(this.messageStore instanceof SimpleMessageStore)) {
-				this.messageStore.removeMessagesFromGroup(this.messageGroupId, message);
-			}
+
 			handleMessageInternal(message);
 		}
 		else {
@@ -587,28 +579,6 @@ private void doReleaseMessage(Message<?> message) {
 		}
 	}
 
-	private boolean removeDelayedMessageFromMessageStore(Message<?> message) {
-		if (this.messageStore instanceof SimpleMessageStore) {
-			this.removeReleasedMessageLock.lock();
-			try {
-				Collection<Message<?>> messages = this.messageStore.getMessageGroup(this.messageGroupId).getMessages();
-				if (messages.contains(message)) {
-					this.messageStore.removeMessagesFromGroup(this.messageGroupId, message);
-					return true;
-				}
-				else {
-					return false;
-				}
-			}
-			finally {
-				this.removeReleasedMessageLock.unlock();
-			}
-		}
-		else {
-			return ((MessageStore) this.messageStore).removeMessage(message.getHeaders().getId()) != null;
-		}
-	}
-
 	@Override
 	public int getDelayedMessageCount() {
 		return this.messageStore.messageGroupSize(this.messageGroupId);
@@ -628,17 +598,17 @@ public void reschedulePersistedMessages() {
 			try (Stream<Message<?>> messageStream = messageGroup.streamMessages()) {
 				TaskScheduler taskScheduler = getTaskScheduler();
 				messageStream.forEach((message) -> // NOSONAR
-				taskScheduler.schedule(() -> {
-					// This is fine to keep the reference to the message,
-					// because the scheduled task is performed immediately.
-					long delay = determineDelayForMessage(message);
-					if (delay > 0) {
-						releaseMessageAfterDelay(message, delay);
-					}
-					else {
-						releaseMessage(message);
-					}
-				}, Instant.now()));
+						taskScheduler.schedule(() -> {
+							// This is fine to keep the reference to the message,
+							// because the scheduled task is performed immediately.
+							long delay = determineDelayForMessage(message);
+							if (delay > 0) {
+								releaseMessageAfterDelay(message, delay);
+							}
+							else {
+								releaseMessage(message);
+							}
+						}, Instant.now()));
 			}
 		}
 		finally {

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

@@ -268,6 +268,42 @@ public void removeMessagesFromGroup(Object groupId, Collection<Message<?>> messa
 		}
 	}
 
+	@Override
+	@Nullable
+	public Message<?> getMessageFromGroup(Object groupId, UUID messageId) {
+		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);
+		Assert.notNull(messageId, "'messageId' must not be null");
+		Object object = doRetrieve(this.messagePrefix + groupId + '_' + messageId);
+		if (object != null) {
+			return extractMessage(object);
+		}
+		else {
+			return null;
+		}
+	}
+
+	@Override
+	public boolean removeMessageFromGroupById(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);
+		if (mgm != null) {
+			Assert.isInstanceOf(MessageGroupMetadata.class, mgm);
+			MessageGroupMetadata messageGroupMetadata = (MessageGroupMetadata) mgm;
+
+			if (messageGroupMetadata.getMessageIds().contains(messageId)) {
+				messageGroupMetadata.remove(messageId);
+				String groupToMessageId = this.messagePrefix + groupId + '_' + messageId;
+				if (doRemove(groupToMessageId) != null) {
+					messageGroupMetadata.setLastModified(System.currentTimeMillis());
+					doStore(this.groupPrefix + groupId, messageGroupMetadata);
+					return true;
+				}
+			}
+		}
+		return false;
+	}
+
 	@Override
 	public void completeGroup(Object groupId) {
 		Assert.notNull(groupId, GROUP_ID_MUST_NOT_BE_NULL);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -18,10 +18,12 @@
 
 import java.util.Collection;
 import java.util.Iterator;
+import java.util.UUID;
 import java.util.stream.Stream;
 
 import org.springframework.jmx.export.annotation.ManagedAttribute;
 import org.springframework.jmx.export.annotation.ManagedOperation;
+import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
 
 /**
@@ -71,6 +73,30 @@ public interface MessageGroupStore extends BasicMessageGroupStore {
 	 */
 	void removeMessagesFromGroup(Object key, Message<?>... messages);
 
+	/**
+	 * Retrieve a {@link Message} from a group by id.
+	 * Return {@code null} if message does not belong to the requested group.
+	 * @param groupId The groupId for the group containing the message.
+	 * @param messageId The message id.
+	 * @return message by id if it belongs to requested group.
+	 * @since 6.1.5
+	 */
+	@Nullable
+	default Message<?> getMessageFromGroup(Object groupId, UUID messageId) {
+		throw new UnsupportedOperationException("Not supported for this store");
+	}
+
+	/**
+	 * Deletion the message from the group.
+	 * @param groupId The groupId for the group containing the message.
+	 * @param messageId The message id to be removed.
+	 * @return true if message has been removed.
+	 * @since 6.1.5
+	 */
+	default boolean removeMessageFromGroupById(Object groupId, UUID messageId) {
+		throw new UnsupportedOperationException("Not supported for this store");
+	}
+
 	/**
 	 * 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.
@@ -114,7 +140,7 @@ public interface MessageGroupStore extends BasicMessageGroupStore {
 
 	/**
 	 * Completes this MessageGroup. Completion of the MessageGroup generally means
-	 * that this group should not be allowing any more mutating operation to be performed on it.
+	 * that this group should not be allowing anymore mutating operation to be performed on it.
 	 * For example any attempt to add/remove new Message form the group should not be allowed.
 	 * @param groupId The group identifier.
 	 */

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

@@ -28,6 +28,7 @@
 import org.springframework.integration.support.locks.LockRegistry;
 import org.springframework.integration.util.UpperBound;
 import org.springframework.jmx.export.annotation.ManagedAttribute;
+import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessagingException;
 import org.springframework.util.Assert;
@@ -382,6 +383,52 @@ public void removeMessagesFromGroup(Object groupId, Collection<Message<?>> messa
 		}
 	}
 
+	@Override
+	@Nullable
+	public Message<?> getMessageFromGroup(Object groupId, UUID messageId) {
+		MessageGroup group = this.groupIdToMessageGroup.get(groupId);
+		Assert.notNull(group,
+				() -> MESSAGE_GROUP_FOR_GROUP_ID + groupId + "' does not exists");
+		for (Message<?> message : group.getMessages()) {
+			if (messageId.equals(message.getHeaders().getId())) {
+				return message;
+			}
+		}
+		return null;
+	}
+
+	@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();
+			}
+		}
+		catch (InterruptedException ex) {
+			Thread.currentThread().interrupt();
+			throw new MessagingException(INTERRUPTED_WHILE_OBTAINING_LOCK, ex);
+		}
+	}
+
 	@Override
 	public Iterator<MessageGroup> iterator() {
 		return new HashSet<>(this.groupIdToMessageGroup.values()).iterator();

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

@@ -54,6 +54,7 @@
 import org.springframework.jdbc.support.lob.DefaultLobHandler;
 import org.springframework.jdbc.support.lob.LobHandler;
 import org.springframework.jmx.export.annotation.ManagedAttribute;
+import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
@@ -163,6 +164,14 @@ SELECT min(CREATED_DATE)
 				where MESSAGE_ID=? and REGION=?
 				"""),
 
+		GET_MESSAGE_FROM_GROUP("""
+				SELECT m.MESSAGE_ID, m.CREATED_DATE, m.MESSAGE_BYTES
+				from %PREFIX%MESSAGE m
+				inner join %PREFIX%GROUP_TO_MESSAGE gm
+					on m.MESSAGE_ID = gm.MESSAGE_ID
+				where gm.MESSAGE_ID=? and gm.GROUP_KEY = ? and gm.REGION=?
+				"""),
+
 		GET_MESSAGE_COUNT("""
 				SELECT COUNT(MESSAGE_ID)
 				from %PREFIX%MESSAGE
@@ -601,6 +610,31 @@ public void removeMessagesFromGroup(Object groupId, Collection<Message<?>> messa
 		updateMessageGroup(groupKey);
 	}
 
+	@Override
+	@Nullable
+	public Message<?> getMessageFromGroup(Object groupId, UUID messageId) {
+		List<Message<?>> list =
+				this.jdbcTemplate.query(getQuery(Query.GET_MESSAGE_FROM_GROUP), this.mapper,
+						getKey(messageId), getKey(groupId), this.region);
+		if (list.isEmpty()) {
+			return null;
+		}
+		return list.get(0);
+	}
+
+	@Override
+	public boolean removeMessageFromGroupById(Object groupId, UUID messageId) {
+		String groupKey = getKey(groupId);
+		String messageKey = getKey(messageId);
+		int messageToGroupRemoved =
+				this.jdbcTemplate.update(getQuery(Query.REMOVE_MESSAGE_FROM_GROUP), groupKey, messageKey, this.region);
+		if (messageToGroupRemoved > 0) {
+			return this.jdbcTemplate.update(getQuery(Query.DELETE_MESSAGE),
+					messageKey, this.region, messageKey, this.region) > 0;
+		}
+		return false;
+	}
+
 	@Override
 	public void removeMessageGroup(Object groupId) {
 		String groupKey = getKey(groupId);

spring-integration-jdbc/src/test/java/org/springframework/integration/jdbc/DelayerHandlerRescheduleIntegrationTests-context.xml

@@ -23,7 +23,7 @@
 	<delayer id="#{T (org.springframework.integration.jdbc.DelayerHandlerRescheduleIntegrationTests).DELAYER_ID}"
 			 input-channel="input"
 			 output-channel="output"
-			 default-delay="10000"
+			 default-delay="1000"
 			 message-store="messageStore"/>
 
 	<channel id="transactionalDelayerOutput"/>

spring-integration-jdbc/src/test/java/org/springframework/integration/jdbc/DelayerHandlerRescheduleIntegrationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -31,7 +31,6 @@
 import org.springframework.integration.store.MessageGroupStore;
 import org.springframework.integration.store.SimpleMessageGroup;
 import org.springframework.integration.support.MessageBuilder;
-import org.springframework.integration.test.condition.LongRunningTest;
 import org.springframework.integration.test.util.TestUtils;
 import org.springframework.integration.util.UUIDConverter;
 import org.springframework.jdbc.datasource.embedded.EmbeddedDatabase;
@@ -47,13 +46,12 @@
 import org.springframework.transaction.support.TransactionSynchronizationManager;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.fail;
+import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
 
 /**
  * @author Artem Bilan
  * @author Gary Russell
  */
-@LongRunningTest
 public class DelayerHandlerRescheduleIntegrationTests {
 
 	public static final String DELAYER_ID = "delayerWithJdbcMS";
@@ -98,15 +96,9 @@ public void testDelayerHandlerRescheduleWithJdbcMessageStore() throws Exception
 		taskScheduler.getScheduledExecutor().awaitTermination(10, TimeUnit.SECONDS);
 		context.close();
 
-		try {
-			context.getBean("input", MessageChannel.class);
-			fail("IllegalStateException expected");
-		}
-		catch (Exception e) {
-			assertThat(e instanceof IllegalStateException).isTrue();
-			assertThat(e.getMessage().contains("BeanFactory not initialized or already closed - call 'refresh'"))
-					.isTrue();
-		}
+		assertThatIllegalStateException()
+				.isThrownBy(() -> context.getBean("input", MessageChannel.class))
+				.withMessageContaining("BeanFactory not initialized or already closed - call 'refresh'");
 
 		String delayerMessageGroupId = UUIDConverter.getUUID(DELAYER_ID + ".messageGroupId").toString();