GH-3446: Stream support in the MessageGroupStore

Fixes #3446

* For better resources utilization provide a `Stream<Message<?>>` API
on the `MessageGroupStore`, `MessageGroup` and `MessageGroupQueue`
* Use this API in the `DelayHandler` when it reschedules persisted messages

Author: artembilan

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -26,6 +26,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Stream;
 
 import org.aopalliance.aop.Advice;
 
@@ -321,7 +322,7 @@ protected boolean shouldCopyRequestHeaders() {
 	}
 
 	/**
-	 * Checks if 'requestMessage' wasn't delayed before ({@link #releaseMessageAfterDelay}
+	 * Check if 'requestMessage' wasn't delayed before ({@link #releaseMessageAfterDelay}
 	 * and {@link DelayHandler.DelayedMessageWrapper}). Than determine 'delay' for
 	 * 'requestMessage' ({@link #determineDelayForMessage}) and if {@code delay > 0}
 	 * schedules 'releaseMessage' task after 'delay'.
@@ -562,9 +563,10 @@ public int getDelayedMessageCount() {
 	@Override
 	public synchronized void reschedulePersistedMessages() {
 		MessageGroup messageGroup = this.messageStore.getMessageGroup(this.messageGroupId);
-		for (final Message<?> message : messageGroup.getMessages()) {
-			getTaskScheduler()
-					.schedule(() -> {
+		try (Stream<Message<?>> messageStream = messageGroup.streamMessages()) {
+			TaskScheduler taskScheduler = getTaskScheduler();
+			messageStream.forEach((message) ->
+					taskScheduler.schedule(() -> {
 						// This is fine to keep the reference to the message,
 						// because the scheduled task is performed immediately.
 						long delay = determineDelayForMessage(message);
@@ -574,12 +576,12 @@ public synchronized void reschedulePersistedMessages() {
 						else {
 							releaseMessage(message);
 						}
-					}, new Date());
+					}, new Date()));
 		}
 	}
 
 	/**
-	 * Handles {@link ContextRefreshedEvent} to invoke
+	 * Handle {@link ContextRefreshedEvent} to invoke
 	 * {@link #reschedulePersistedMessages} as late as possible after application context
 	 * startup. Also it checks {@link #initialized} to ignore other
 	 * {@link ContextRefreshedEvent}s which may be published in the 'parent-child'

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -24,6 +24,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.springframework.jmx.export.annotation.ManagedAttribute;
 import org.springframework.messaging.Message;
@@ -160,7 +161,7 @@ public Message<?> removeMessage(UUID id) {
 	@Override
 	@ManagedAttribute
 	public long getMessageCount() {
-		Collection<?> messageIds = doListKeys(this.messagePrefix + "*");
+		Collection<?> messageIds = doListKeys(this.messagePrefix + '*');
 		return (messageIds != null) ? messageIds.size() : 0;
 	}
 
@@ -346,11 +347,19 @@ public Collection<Message<?>> getMessagesForGroup(Object groupId) {
 		return messages;
 	}
 
+	@Override
+	public Stream<Message<?>> streamMessagesForGroup(Object groupId) {
+		return getGroupMetadata(groupId)
+				.getMessageIds()
+				.stream()
+				.map(this::getMessage);
+	}
+
 	@Override
 	@SuppressWarnings("unchecked")
 	public Iterator<MessageGroup> iterator() {
 		final Iterator<?> idIterator = normalizeKeys(
-				(Collection<String>) doListKeys(this.groupPrefix + "*"))
+				(Collection<String>) doListKeys(this.groupPrefix + '*'))
 				.iterator();
 		return new MessageGroupIterator(idIterator);
 	}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -17,14 +17,15 @@
 package org.springframework.integration.store;
 
 import java.util.Collection;
+import java.util.stream.Stream;
 
 import org.springframework.messaging.Message;
 
 /**
  * A group of messages that are correlated with each other and should be processed in the same context.
  * <p>
- * The message group allows implementations to be mutable, but this behavior is optional. Implementations should take
- * care to document their thread safety and mutability.
+ * The message group allows implementations to be mutable, but this behavior is optional.
+ * Implementations should take care to document their thread safety and mutability.
  *
  * @author Dave Syer
  * @author Oleg Zhurakousky
@@ -35,7 +36,6 @@ public interface MessageGroup {
 
 	/**
 	 * Query if the message can be added.
-	 *
 	 * @param message The message.
 	 * @return true if the message can be added.
 	 */
@@ -57,12 +57,20 @@ public interface MessageGroup {
 	boolean remove(Message<?> messageToRemove);
 
 	/**
-	 * Returns all available Messages from the group at the time of invocation
-	 *
+	 * Return all available Messages from the group at the time of invocation
 	 * @return The messages.
 	 */
 	Collection<Message<?>> getMessages();
 
+	/**
+	 * Return a stream for messages stored in this group.
+	 * @return the {@link Stream} for messages in this group.
+	 * @since 5.5
+	 */
+	default Stream<Message<?>> streamMessages() {
+		return getMessages().stream();
+	}
+
 	/**
 	 * @return the key that links these messages together
 	 */

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -25,6 +25,7 @@
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.stream.Stream;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -42,6 +43,7 @@
  * @author Oleg Zhurakousky
  * @author Gunnar Hillert
  * @author Gary Russell
+ * @author Artem Bilan
  *
  * @since 2.0
  *
@@ -118,7 +120,7 @@ public void setPriority(boolean priority) {
 
 	@Override
 	public Iterator<Message<?>> iterator() {
-		return getMessages().iterator();
+		return stream().iterator();
 	}
 
 	/**
@@ -164,29 +166,24 @@ public int size() {
 
 	@Override
 	public Message<?> peek() {
-		Message<?> message = null;
-		final Lock lock = this.storeLock;
 		try {
-			lock.lockInterruptibly();
-			try {
-				Collection<Message<?>> messages = getMessages();
-				if (!messages.isEmpty()) {
-					message = messages.iterator().next();
-				}
+			this.storeLock.lockInterruptibly();
+			try (Stream<Message<?>> messageStream = stream()) {
+				return messageStream.findFirst().orElse(null);
 			}
 			finally {
-				lock.unlock();
+				this.storeLock.unlock();
 			}
 		}
 		catch (InterruptedException e) {
 			Thread.currentThread().interrupt();
 		}
-		return message;
+		return null;
 	}
 
 	@Override
 	public Message<?> poll(long timeout, TimeUnit unit) throws InterruptedException {
-		Message<?> message = null;
+		Message<?> message;
 		long timeoutInNanos = unit.toNanos(timeout);
 		final Lock lock = this.storeLock;
 		lock.lockInterruptibly();
@@ -325,7 +322,7 @@ public int remainingCapacity() {
 
 	@Override
 	public Message<?> take() throws InterruptedException {
-		Message<?> message = null;
+		Message<?> message;
 		final Lock lock = this.storeLock;
 		lock.lockInterruptibly();
 
@@ -346,6 +343,11 @@ protected Collection<Message<?>> getMessages() {
 		return this.messageGroupStore.getMessageGroup(this.groupId).getMessages();
 	}
 
+	@Override
+	public Stream<Message<?>> stream() {
+		return this.messageGroupStore.getMessageGroup(this.groupId).streamMessages();
+	}
+
 	/**
 	 * It is assumed that the 'storeLock' is being held by the caller, otherwise
 	 * IllegalMonitorStateException may be thrown

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2021 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,6 +18,7 @@
 
 import java.util.Collection;
 import java.util.Iterator;
+import java.util.stream.Stream;
 
 import org.springframework.jmx.export.annotation.ManagedAttribute;
 import org.springframework.jmx.export.annotation.ManagedOperation;
@@ -142,6 +143,16 @@ public interface MessageGroupStore extends BasicMessageGroupStore {
 	 */
 	Collection<Message<?>> getMessagesForGroup(Object groupId);
 
+	/**
+	 * Return a stream for messages stored in the provided group.
+	 * @param groupId the group id to retrieve messages.
+	 * @return the {@link Stream} for messages in this group.
+	 * @since 5.5
+	 */
+	default Stream<Message<?>> streamMessagesForGroup(Object groupId) {
+		return getMessagesForGroup(groupId).stream();
+	}
+
 	/**
 	 * Invoked when a MessageGroupStore expires a group.
 	 */

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2020 the original author or authors.
+ * Copyright 2016-2021 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.
@@ -20,6 +20,8 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.Spliterator;
+import java.util.stream.Stream;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -29,6 +31,7 @@
 
 /**
  * @author Artem Bilan
+ *
  * @since 4.3
  */
 class PersistentMessageGroup implements MessageGroup {
@@ -59,6 +62,16 @@ public Collection<Message<?>> getMessages() {
 		return Collections.unmodifiableCollection(this.messages);
 	}
 
+	/**
+	 * The resulting {@link Stream} must be closed after use,
+	 * it can be declared as a resource in a {@code try-with-resources} statement.
+	 * @return the stream of messages in this group.
+	 */
+	@Override
+	public Stream<Message<?>> streamMessages() {
+		return this.messageGroupStore.streamMessagesForGroup(this.original.getGroupId());
+	}
+
 	@Override
 	public Message<?> getOne() {
 		if (this.oneMessage == null) {
@@ -224,6 +237,16 @@ public int size() {
 			return PersistentMessageGroup.this.size();
 		}
 
+		@Override
+		public Spliterator<Message<?>> spliterator() {
+			return streamMessages().spliterator();
+		}
+
+		@Override
+		public Stream<Message<?>> stream() {
+			return streamMessages();
+		}
+
 	}
 
 }

spring-integration-core/src/test/java/org/springframework/integration/config/xml/DelayerParserTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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,8 +22,7 @@
 import java.util.HashMap;
 import java.util.List;
 
-import org.junit.Test;
-import org.junit.runner.RunWith;
+import org.junit.jupiter.api.Test;
 
 import org.springframework.beans.DirectFieldAccessor;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -35,8 +34,7 @@
 import org.springframework.integration.test.util.TestUtils;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageHandler;
-import org.springframework.test.context.ContextConfiguration;
-import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
 import org.springframework.transaction.TransactionDefinition;
 import org.springframework.transaction.interceptor.MatchAlwaysTransactionAttributeSource;
 import org.springframework.transaction.interceptor.NameMatchTransactionAttributeSource;
@@ -51,8 +49,7 @@
  *
  * @since 1.0.3
  */
-@RunWith(SpringJUnit4ClassRunner.class)
-@ContextConfiguration
+@SpringJUnitConfig
 public class DelayerParserTests {
 
 	@Autowired
@@ -70,7 +67,7 @@ public void defaultScheduler() {
 				.isEqualTo("headers.foo");
 		assertThat(TestUtils.getPropertyValue(delayHandler, "messagingTemplate.sendTimeout", Long.class))
 				.isEqualTo(987L);
-		assertThat(TestUtils.getPropertyValue(delayHandler, "taskScheduler")).isNull();
+		assertThat(TestUtils.getPropertyValue(delayHandler, "taskScheduler")).isNotNull();
 	}
 
 	@Test
@@ -110,7 +107,8 @@ public void transactionalSubElement() throws Exception {
 		assertThat(adviceChain.size()).isEqualTo(1);
 		Object advice = adviceChain.get(0);
 		assertThat(advice instanceof TransactionInterceptor).isTrue();
-		TransactionAttributeSource transactionAttributeSource = ((TransactionInterceptor) advice).getTransactionAttributeSource();
+		TransactionAttributeSource transactionAttributeSource =
+				((TransactionInterceptor) advice).getTransactionAttributeSource();
 		assertThat(transactionAttributeSource instanceof MatchAlwaysTransactionAttributeSource).isTrue();
 		Method method = MessageHandler.class.getMethod("handleMessage", Message.class);
 		TransactionDefinition definition = transactionAttributeSource.getTransactionAttribute(method, null);
@@ -130,7 +128,8 @@ public void adviceChainSubElement() {
 
 		Object txAdvice = adviceChain.get(1);
 		assertThat(txAdvice.getClass()).isEqualTo(TransactionInterceptor.class);
-		TransactionAttributeSource transactionAttributeSource = ((TransactionInterceptor) txAdvice).getTransactionAttributeSource();
+		TransactionAttributeSource transactionAttributeSource =
+				((TransactionInterceptor) txAdvice).getTransactionAttributeSource();
 		assertThat(transactionAttributeSource.getClass()).isEqualTo(NameMatchTransactionAttributeSource.class);
 		HashMap<?, ?> nameMap = TestUtils.getPropertyValue(transactionAttributeSource, "nameMap", HashMap.class);
 		assertThat(nameMap.toString()).isEqualTo("{*=PROPAGATION_REQUIRES_NEW,ISOLATION_DEFAULT,readOnly}");