Handle STOMP messages from client in order

See gh-21798

Author: rstoyanchev

spring-messaging/src/main/java/org/springframework/messaging/simp/broker/OrderedMessageChannelDecorator.java

@@ -113,7 +113,12 @@ private void sendNextMessage() {
 		}
 	}
 
+	/**
+	 * Remove the message from the top of the queue, but only if it matches,
+	 * i.e. hasn't been removed already.
+	 */
 	private boolean removeMessage(Message<?> message) {
+		// Remove only if not removed already
 		Message<?> next = this.messages.peek();
 		if (next == message) {
 			this.messages.remove();
@@ -181,7 +186,7 @@ private PostHandleTask(Message<?> message) {
 		@Override
 		public void run() {
 			if (this.handledCount == null || this.handledCount.addAndGet(1) == subscriberCount) {
-				if (OrderedMessageChannelDecorator.this.removeMessage(message)) {
+				if (OrderedMessageChannelDecorator.this.removeMessage(this.message)) {
 					sendNextMessage();
 				}
 			}

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/StompEndpointRegistry.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 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.
@@ -52,4 +52,18 @@ public interface StompEndpointRegistry {
 	 */
 	WebMvcStompEndpointRegistry setErrorHandler(StompSubProtocolErrorHandler errorHandler);
 
+	/**
+	 * Whether to handle client messages sequentially in the order in which
+	 * they were received.
+	 * <p>By default messages sent to the {@code "clientInboundChannel"} may
+	 * be handled in parallel and not in the same order as they were received
+	 * because the channel is backed by a ThreadPoolExecutor that in turn does
+	 * not guarantee processing in order.
+	 * <p>When this flag is set to {@code true} messages within the same session
+	 * will be sent to the {@code "clientInboundChannel"} one at a time in
+	 * order to preserve the order in which they were received.
+	 * @since 6.1
+	 */
+	WebMvcStompEndpointRegistry setPreserveReceiveOrder(boolean preserveReceiveOrder);
+
 }

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebMvcStompEndpointRegistry.java

@@ -142,6 +142,15 @@ public WebMvcStompEndpointRegistry setErrorHandler(StompSubProtocolErrorHandler
 		return this;
 	}
 
+	public WebMvcStompEndpointRegistry setPreserveReceiveOrder(boolean preserveReceiveOrder) {
+		this.stompHandler.setPreserveReceiveOrder(preserveReceiveOrder);
+		return this;
+	}
+
+	protected boolean isPreserveReceiveOrder() {
+		return this.stompHandler.isPreserveReceiveOrder();
+	}
+
 	protected void setApplicationContext(ApplicationContext applicationContext) {
 		this.stompHandler.setApplicationEventPublisher(applicationContext);
 	}

spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java

@@ -28,6 +28,7 @@
 import org.springframework.messaging.simp.SimpSessionScope;
 import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler;
 import org.springframework.messaging.simp.broker.AbstractBrokerMessageHandler;
+import org.springframework.messaging.simp.broker.OrderedMessageChannelDecorator;
 import org.springframework.messaging.simp.config.AbstractMessageBrokerConfiguration;
 import org.springframework.messaging.simp.stomp.StompBrokerRelayMessageHandler;
 import org.springframework.messaging.simp.user.SimpUserRegistry;
@@ -80,7 +81,8 @@ protected SimpUserRegistry createLocalUserRegistry(@Nullable Integer order) {
 
 	@Bean
 	public HandlerMapping stompWebSocketHandlerMapping(
-			WebSocketHandler subProtocolWebSocketHandler, TaskScheduler messageBrokerTaskScheduler) {
+			WebSocketHandler subProtocolWebSocketHandler, TaskScheduler messageBrokerTaskScheduler,
+			AbstractSubscribableChannel clientInboundChannel) {
 
 		WebSocketHandler handler = decorateWebSocketHandler(subProtocolWebSocketHandler);
 		WebMvcStompEndpointRegistry registry =
@@ -90,6 +92,7 @@ public HandlerMapping stompWebSocketHandlerMapping(
 			registry.setApplicationContext(applicationContext);
 		}
 		registerStompEndpoints(registry);
+		OrderedMessageChannelDecorator.configureInterceptor(clientInboundChannel, registry.isPreserveReceiveOrder());
 		return registry.getHandlerMapping();
 	}
 

spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java

@@ -108,6 +108,10 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 	@Nullable
 	private MessageHeaderInitializer headerInitializer;
 
+	private boolean preserveReceiveOrder;
+
+	private final Map<String, MessageChannel> messageChannels = new ConcurrentHashMap<>();
+
 	private final Map<String, Principal> stompAuthentications = new ConcurrentHashMap<>();
 
 	@Nullable
@@ -193,6 +197,30 @@ public MessageHeaderInitializer getHeaderInitializer() {
 		return this.headerInitializer;
 	}
 
+	/**
+	 * Whether client messages must be handled in the order received.
+	 * <p>By default messages sent to the {@code "clientInboundChannel"} may
+	 * not be handled in the same order because the channel is backed by a
+	 * ThreadPoolExecutor that in turn does not guarantee processing in order.
+	 * <p>When this flag is set to {@code true} messages within the same session
+	 * will be sent to the {@code "clientInboundChannel"} one at a time to
+	 * preserve the order in which they were received.
+	 * @param preserveReceiveOrder whether to publish in order
+	 * @since 6.1
+	 */
+	public void setPreserveReceiveOrder(boolean preserveReceiveOrder) {
+		this.preserveReceiveOrder = preserveReceiveOrder;
+	}
+
+	/**
+	 * Whether the handler is configured to handle inbound messages in the
+	 * order in which they were received.
+	 * @since 6.1
+	 */
+	public boolean isPreserveReceiveOrder() {
+		return this.preserveReceiveOrder;
+	}
+
 	@Override
 	public List<String> getSupportedProtocols() {
 		return Arrays.asList("v10.stomp", "v11.stomp", "v12.stomp");
@@ -268,6 +296,12 @@ else if (webSocketMessage instanceof BinaryMessage binaryMessage) {
 			return;
 		}
 
+		MessageChannel channelToUse =
+				(this.messageChannels.computeIfAbsent(session.getId(),
+						id -> this.preserveReceiveOrder ?
+								new OrderedMessageChannelDecorator(outputChannel, logger) :
+								outputChannel));
+
 		for (Message<byte[]> message : messages) {
 			StompHeaderAccessor headerAccessor =
 					MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
@@ -307,7 +341,7 @@ else if (StompCommand.DISCONNECT.equals(command)) {
 
 				try {
 					SimpAttributesContextHolder.setAttributesFromMessage(message);
-					sent = outputChannel.send(message);
+					sent = channelToUse.send(message);
 
 					if (sent) {
 						if (this.eventPublisher != null) {
@@ -652,6 +686,7 @@ public void afterSessionEnded(WebSocketSession session, CloseStatus closeStatus,
 			outputChannel.send(message);
 		}
 		finally {
+			this.messageChannels.remove(session.getId());
 			this.stompAuthentications.remove(session.getId());
 			SimpAttributesContextHolder.resetAttributes();
 			simpAttributes.sessionCompleted();