Support user destinations with multiple app servers

This change adds support for broadcasting messages with unresolved
user destinations so that other servers can try to resolve it.
That enables sending messages to users who may be connected to a
different server.

Issue: SPR-11620

spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageHeaderAccessor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2015 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.
@@ -69,14 +69,17 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor {
 
 
 	/**
-	 * For internal use.
-	 * <p>The original destination used by a client when subscribing. Such a
-	 * destination may have been modified (e.g. user destinations) on the server
-	 * side. This header provides a hint so messages sent to clients may have
-	 * a destination matching to their original subscription.
+	 * A header for internal use with "user" destinations where we need to
+	 * restore the destination prior to sending messages to clients.
 	 */
 	public static final String ORIGINAL_DESTINATION = "simpOrigDestination";
 
+	/**
+	 * A header that indicates to the broker that the sender will ignore errors.
+	 * The header is simply checked for presence or absence.
+	 */
+	public static final String IGNORE_ERROR = "simpIgnoreError";
+
 
 	/**
 	 * A constructor for creating new message headers.

spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java

@@ -17,7 +17,10 @@
 package org.springframework.messaging.simp.config;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.factory.BeanInitializationException;
@@ -25,6 +28,7 @@ import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
 import org.springframework.context.annotation.Bean;
 import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHandler;
 import org.springframework.messaging.converter.ByteArrayMessageConverter;
 import org.springframework.messaging.converter.CompositeMessageConverter;
 import org.springframework.messaging.converter.DefaultContentTypeResolver;
@@ -37,6 +41,7 @@ import org.springframework.messaging.simp.SimpMessagingTemplate;
 import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler;
 import org.springframework.messaging.simp.broker.AbstractBrokerMessageHandler;
 import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
+import org.springframework.messaging.simp.stomp.StompBrokerRelayMessageHandler;
 import org.springframework.messaging.simp.user.DefaultUserDestinationResolver;
 import org.springframework.messaging.simp.user.DefaultUserSessionRegistry;
 import org.springframework.messaging.simp.user.UserDestinationMessageHandler;
@@ -278,13 +283,26 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC
 
 	@Bean
 	public AbstractBrokerMessageHandler stompBrokerRelayMessageHandler() {
-		AbstractBrokerMessageHandler handler = getBrokerRegistry().getStompBrokerRelay(brokerChannel());
-		return (handler != null ? handler : new NoOpBrokerMessageHandler());
+		StompBrokerRelayMessageHandler handler = getBrokerRegistry().getStompBrokerRelay(brokerChannel());
+		if (handler == null) {
+			return new NoOpBrokerMessageHandler();
+		}
+		String destination = getBrokerRegistry().getUserDestinationBroadcast();
+		if (destination != null) {
+			Map<String, MessageHandler> map = new HashMap<String, MessageHandler>(1);
+			map.put(destination, userDestinationMessageHandler());
+			handler.setSystemSubscriptions(map);
+		}
+		return handler;
 	}
 
 	@Bean
 	public UserDestinationMessageHandler userDestinationMessageHandler() {
-		return new UserDestinationMessageHandler(clientInboundChannel(), brokerChannel(), userDestinationResolver());
+		UserDestinationMessageHandler handler = new UserDestinationMessageHandler(clientInboundChannel(),
+				brokerChannel(), userDestinationResolver());
+		String destination = getBrokerRegistry().getUserDestinationBroadcast();
+		handler.setUserDestinationBroadcast(destination);
+		return handler;
 	}
 
 	@Bean

spring-messaging/src/main/java/org/springframework/messaging/simp/config/MessageBrokerRegistry.java

@@ -49,6 +49,8 @@ public class MessageBrokerRegistry {
 
 	private String userDestinationPrefix;
 
+	private String userDestinationBroadcast;
+
 	private PathMatcher pathMatcher;
 
 
@@ -137,6 +139,24 @@ public class MessageBrokerRegistry {
 		return this.userDestinationPrefix;
 	}
 
+	/**
+	 * Set a destination to broadcast messages to that remain unresolved because
+	 * the user is not connected. In a multi-application server scenario this
+	 * gives other application servers a chance to try.
+	 * <p><strong>Note:</strong> this option applies only when the
+	 * {@link #enableStompBrokerRelay "broker relay"} is enabled.
+	 * <p>By default this is not set.
+	 * @param destination the destination to forward unresolved
+	 * messages to, e.g. "/topic/unresolved-user-destination".
+	 */
+	public void setUserDestinationBroadcast(String destination) {
+		this.userDestinationBroadcast = destination;
+	}
+
+	protected String getUserDestinationBroadcast() {
+		return this.userDestinationBroadcast;
+	}
+
 	/**
 	 * Configure the PathMatcher to use to match the destinations of incoming
 	 * messages to {@code @MessageMapping} and {@code @SubscribeMapping} methods.

spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java

@@ -17,6 +17,7 @@
 package org.springframework.messaging.simp.stomp;
 
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
@@ -26,6 +27,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.MessageDeliveryException;
+import org.springframework.messaging.MessageHandler;
 import org.springframework.messaging.MessageHeaders;
 import org.springframework.messaging.SubscribableChannel;
 import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
@@ -112,6 +114,8 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 
 	private String virtualHost;
 
+	private final Map<String, MessageHandler> systemSubscriptions = new HashMap<String, MessageHandler>(4);
+
 	private TcpOperations<byte[]> tcpClient;
 
 	private MessageHeaderInitializer headerInitializer;
@@ -281,6 +285,27 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 		return this.systemPasscode;
 	}
 
+	/**
+	 * Configure one more destinations to subscribe to on the shared "system"
+	 * connection along with MessageHandler's to handle received messages.
+	 * <p>This is for internal use in a multi-application server scenario where
+	 * servers forward messages to each other (e.g. unresolved user destinations).
+	 * @param subscriptions the destinations to subscribe to.
+	 */
+	public void setSystemSubscriptions(Map<String, MessageHandler> subscriptions) {
+		this.systemSubscriptions.clear();
+		if (subscriptions != null) {
+			this.systemSubscriptions.putAll(subscriptions);
+		}
+	}
+
+	/**
+	 * Return the configured map with subscriptions on the "system" connection.
+	 */
+	public Map<String, MessageHandler> getSystemSubscriptions() {
+		return this.systemSubscriptions;
+	}
+
 	/**
 	 * Set the value of the "host" header to use in STOMP CONNECT frames. When this
 	 * property is configured, a "host" header will be added to every STOMP frame sent to
@@ -532,6 +557,10 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 			return this.sessionId;
 		}
 
+		protected TcpConnection<byte[]> getTcpConnection() {
+			return this.tcpConnection;
+		}
+
 		@Override
 		public void afterConnected(TcpConnection<byte[]> connection) {
 			if (logger.isDebugEnabled()) {
@@ -579,13 +608,14 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 				headerAccessor.setUser(this.connectHeaders.getUser());
 				headerAccessor.setMessage(errorText);
 				Message<?> errorMessage = MessageBuilder.createMessage(EMPTY_PAYLOAD, headerAccessor.getMessageHeaders());
-				headerAccessor.setImmutable();
-				sendMessageToClient(errorMessage);
+				handleInboundMessage(errorMessage);
 			}
 		}
 
-		protected void sendMessageToClient(Message<?> message) {
+		protected void handleInboundMessage(Message<?> message) {
 			if (this.isRemoteClientSession) {
+				StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
+				accessor.setImmutable();
 				StompBrokerRelayMessageHandler.this.getClientOutboundChannel().send(message);
 			}
 		}
@@ -610,8 +640,7 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 				logger.trace("Received " + accessor.getDetailedLogMessage(message.getPayload()));
 			}
 
-			accessor.setImmutable();
-			sendMessageToClient(message);
+			handleInboundMessage(message);
 		}
 
 		/**
@@ -825,7 +854,6 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 		}
 	}
 
-
 	private class SystemStompConnectionHandler extends StompConnectionHandler {
 
 		public SystemStompConnectionHandler(StompHeaderAccessor connectHeaders) {
@@ -839,6 +867,63 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 			}
 			super.afterStompConnected(connectedHeaders);
 			publishBrokerAvailableEvent();
+			sendSystemSubscriptions();
+		}
+
+		private void sendSystemSubscriptions() {
+			int i = 0;
+			for (String destination : getSystemSubscriptions().keySet()) {
+				StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.SUBSCRIBE);
+				accessor.setSubscriptionId(String.valueOf(i++));
+				accessor.setDestination(destination);
+				if (logger.isDebugEnabled()) {
+					logger.debug("Subscribing to " + destination + " on \"system\" connection.");
+				}
+				TcpConnection<byte[]> conn = getTcpConnection();
+				if (conn != null) {
+					MessageHeaders headers = accessor.getMessageHeaders();
+					conn.send(MessageBuilder.createMessage(EMPTY_PAYLOAD, headers)).addCallback(
+							new ListenableFutureCallback<Void>() {
+								public void onSuccess(Void result) {
+								}
+								public void onFailure(Throwable ex) {
+									String error = "Failed to subscribe in \"system\" session.";
+									handleTcpConnectionFailure(error, ex);
+								}
+							});
+				}
+			}
+		}
+
+		@Override
+		protected void handleInboundMessage(Message<?> message) {
+			StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
+			if (StompCommand.MESSAGE.equals(accessor.getCommand())) {
+				String destination = accessor.getDestination();
+				if (destination == null) {
+					if (logger.isDebugEnabled()) {
+						logger.debug("Got message on \"system\" connection, with no destination: " +
+								accessor.getDetailedLogMessage(message.getPayload()));
+					}
+					return;
+				}
+				if (!getSystemSubscriptions().containsKey(destination)) {
+					if (logger.isDebugEnabled()) {
+						logger.debug("Got message on \"system\" connection with no handler: " +
+								accessor.getDetailedLogMessage(message.getPayload()));
+					}
+					return;
+				}
+				try {
+					MessageHandler handler = getSystemSubscriptions().get(destination);
+					handler.handleMessage(message);
+				}
+				catch (Throwable ex) {
+					if (logger.isDebugEnabled()) {
+						logger.debug("Error while handling message on \"system\" connection.", ex);
+					}
+				}
+			}
 		}
 
 		@Override
@@ -857,7 +942,9 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 		public ListenableFuture<Void> forward(Message<?> message, StompHeaderAccessor accessor) {
 			try {
 				ListenableFuture<Void> future = super.forward(message, accessor);
-				future.get();
+				if (message.getHeaders().get(SimpMessageHeaderAccessor.IGNORE_ERROR) == null) {
+					future.get();
+				}
 				return future;
 			}
 			catch (Throwable ex) {

spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java

@@ -16,7 +16,11 @@
 
 package org.springframework.messaging.simp.user;
 
-import java.util.Set;
+import static org.springframework.messaging.simp.SimpMessageHeaderAccessor.*;
+import static org.springframework.messaging.support.MessageHeaderAccessor.getAccessor;
+
+import java.util.Arrays;
+import java.util.List;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -24,6 +28,7 @@ import org.apache.commons.logging.LogFactory;
 import org.springframework.context.SmartLifecycle;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.MessageHeaders;
 import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.SubscribableChannel;
 import org.springframework.messaging.core.MessageSendingOperations;
@@ -33,6 +38,7 @@ import org.springframework.messaging.simp.SimpMessagingTemplate;
 import org.springframework.messaging.support.MessageBuilder;
 import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
 
 /**
  * {@code MessageHandler} with support for "user" destinations.
@@ -53,9 +59,11 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec
 
 	private final SubscribableChannel brokerChannel;
 
+	private final UserDestinationResolver destinationResolver;
+
 	private final MessageSendingOperations<String> messagingTemplate;
 
-	private final UserDestinationResolver destinationResolver;
+	private BroadcastHandler broadcastHandler;
 
 	private MessageHeaderInitializer headerInitializer;
 
@@ -93,6 +101,25 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec
 		return this.destinationResolver;
 	}
 
+	/**
+	 * Set a destination to broadcast messages to that remain unresolved because
+	 * the user is not connected. In a multi-application server scenario this
+	 * gives other application servers a chance to try.
+	 * <p>By default this is not set.
+	 * @param destination the target destination.
+	 */
+	public void setUserDestinationBroadcast(String destination) {
+		this.broadcastHandler = (StringUtils.hasText(destination) ?
+				new BroadcastHandler(this.messagingTemplate, destination) : null);
+	}
+
+	/**
+	 * Return the configured destination for unresolved messages.
+	 */
+	public String getUserDestinationBroadcast() {
+		return (this.broadcastHandler != null ? this.broadcastHandler.getBroadcastDestination() : null);
+	}
+
 	/**
 	 * Return the messaging template used to send resolved messages to the
 	 * broker channel.
@@ -164,29 +191,35 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec
 
 	@Override
 	public void handleMessage(Message<?> message) throws MessagingException {
+		if (this.broadcastHandler != null) {
+			message = this.broadcastHandler.preHandle(message);
+			if (message == null) {
+				return;
+			}
+		}
 		UserDestinationResult result = this.destinationResolver.resolveDestination(message);
 		if (result == null) {
 			return;
 		}
-		Set<String> destinations = result.getTargetDestinations();
-		if (destinations.isEmpty()) {
+		if (result.getTargetDestinations().isEmpty()) {
 			if (logger.isTraceEnabled()) {
-				logger.trace("No user destinations found for " + result.getSourceDestination());
+				logger.trace("No active sessions for user destination: " + result.getSourceDestination());
+			}
+			if (this.broadcastHandler != null) {
+				this.broadcastHandler.handleUnresolved(message);
 			}
 			return;
 		}
-		if (SimpMessageType.MESSAGE.equals(SimpMessageHeaderAccessor.getMessageType(message.getHeaders()))) {
-			SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.wrap(message);
-			initHeaders(accessor);
-			String header = SimpMessageHeaderAccessor.ORIGINAL_DESTINATION;
-			accessor.setNativeHeader(header, result.getSubscribeDestination());
-			message = MessageBuilder.createMessage(message.getPayload(), accessor.getMessageHeaders());
+		SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.wrap(message);
+		initHeaders(accessor);
+		accessor.setNativeHeader(ORIGINAL_DESTINATION, result.getSubscribeDestination());
+		accessor.setLeaveMutable(true);
+		message = MessageBuilder.createMessage(message.getPayload(), accessor.getMessageHeaders());
+		if (logger.isTraceEnabled()) {
+			logger.trace("Translated " + result.getSourceDestination() + " -> " + result.getTargetDestinations());
 		}
-		if (logger.isDebugEnabled()) {
-			logger.debug("Translated " + result.getSourceDestination() + " -> " + destinations);
-		}
-		for (String destination : destinations) {
-			this.messagingTemplate.send(destination, message);
+		for (String target : result.getTargetDestinations()) {
+			this.messagingTemplate.send(target, message);
 		}
 	}
 
@@ -201,4 +234,73 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec
 		return "UserDestinationMessageHandler[" + this.destinationResolver + "]";
 	}
 
+
+	/**
+	 * A handler that broadcasts locally unresolved messages to the broker and
+	 * also handles similar broadcasts received from the broker.
+	 */
+	private static class BroadcastHandler {
+
+		private static final List<String> NO_COPY_LIST = Arrays.asList("subscription", "message-id");
+
+
+		private final MessageSendingOperations<String> messagingTemplate;
+
+		private final String broadcastDestination;
+
+
+		public BroadcastHandler(MessageSendingOperations<String> template, String destination) {
+			this.messagingTemplate = template;
+			this.broadcastDestination = destination;
+		}
+
+
+		public String getBroadcastDestination() {
+			return this.broadcastDestination;
+		}
+
+		public Message<?> preHandle(Message<?> message) throws MessagingException {
+			String destination = SimpMessageHeaderAccessor.getDestination(message.getHeaders());
+			if (!getBroadcastDestination().equals(destination)) {
+				return message;
+			}
+			SimpMessageHeaderAccessor accessor = getAccessor(message, SimpMessageHeaderAccessor.class);
+			if (accessor.getSessionId() == null) {
+				// Our own broadcast
+				return null;
+			}
+			destination = accessor.getFirstNativeHeader(ORIGINAL_DESTINATION);
+			if (logger.isTraceEnabled()) {
+				logger.trace("Checking unresolved user destination: " + destination);
+			}
+			SimpMessageHeaderAccessor newAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
+			for (String name : accessor.toNativeHeaderMap().keySet()) {
+				if (NO_COPY_LIST.contains(name)) {
+					continue;
+				}
+				newAccessor.setNativeHeader(name, accessor.getFirstNativeHeader(name));
+			}
+			newAccessor.setDestination(destination);
+			newAccessor.setHeader(SimpMessageHeaderAccessor.IGNORE_ERROR, true); // ensure send doesn't block
+			return MessageBuilder.createMessage(message.getPayload(), newAccessor.getMessageHeaders());
+		}
+
+		public void handleUnresolved(Message<?> message) {
+			MessageHeaders headers = message.getHeaders();
+			if (SimpMessageHeaderAccessor.getFirstNativeHeader(ORIGINAL_DESTINATION, headers) != null) {
+				// Re-broadcast
+				return;
+			}
+			SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.wrap(message);
+			String destination = accessor.getDestination();
+			accessor.setNativeHeader(ORIGINAL_DESTINATION, destination);
+			accessor.setLeaveMutable(true);
+			message = MessageBuilder.createMessage(message.getPayload(), accessor.getMessageHeaders());
+			if (logger.isTraceEnabled()) {
+				logger.trace("Translated " + destination + " -> " + getBroadcastDestination());
+			}
+			this.messagingTemplate.send(getBroadcastDestination(), message);
+		}
+	}
+
 }

spring-messaging/src/main/java/org/springframework/messaging/support/NativeMessageHeaderAccessor.java

@@ -115,7 +115,7 @@ public class NativeMessageHeaderAccessor extends MessageHeaderAccessor {
 	 */
 	public boolean containsNativeHeader(String headerName) {
 		Map<String, List<String>> map = getNativeHeaders();
-		return (map != null ? map.containsKey(headerName) : false);
+		return (map != null && map.containsKey(headerName));
 	}
 
 	/**
@@ -207,4 +207,16 @@ public class NativeMessageHeaderAccessor extends MessageHeaderAccessor {
 		return nativeHeaders.remove(name);
 	}
 
+	@SuppressWarnings("unchecked")
+	public static String getFirstNativeHeader(String headerName, Map<String, Object> headers) {
+		Map<String, List<String>> map = (Map<String, List<String>>) headers.get(NATIVE_HEADERS);
+		if (map != null) {
+			List<String> values = map.get(headerName);
+			if (values != null) {
+				return values.get(0);
+			}
+		}
+		return null;
+	}
+
 }