Use strategy for customizing ID/TIMESTAMP generation

Issue: SPR-11468

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

@@ -27,6 +27,7 @@ import org.springframework.messaging.core.AbstractMessageSendingTemplate;
 import org.springframework.messaging.core.MessagePostProcessor;
 import org.springframework.messaging.support.MessageBuilder;
 import org.springframework.messaging.support.MessageHeaderAccessor;
+import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.messaging.support.NativeMessageHeaderAccessor;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
@@ -48,6 +49,8 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplate<String
 
 	private String userDestinationPrefix = "/user/";
 
+	private MessageHeaderInitializer headerInitializer;
+
 	private volatile long sendTimeout = -1;
 
 
@@ -78,6 +81,23 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplate<String
 		return this.userDestinationPrefix;
 	}
 
+	/**
+	 * Configure a {@link MessageHeaderInitializer} to apply to the headers of all
+	 * messages created through the {@code SimpMessagingTemplate}.
+	 *
+	 * <p>By default this property is not set.
+	 */
+	public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) {
+		this.headerInitializer = headerInitializer;
+	}
+
+	/**
+	 * @return the configured header initializer.
+	 */
+	public MessageHeaderInitializer getHeaderInitializer() {
+		return this.headerInitializer;
+	}
+
 	/**
 	 * @return the messageChannel
 	 */
@@ -146,10 +166,12 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplate<String
 			else {
 				// Try and keep the original accessor type
 				simpAccessor = (SimpMessageHeaderAccessor) MessageHeaderAccessor.getMutableAccessor(message);
+				initHeaders(simpAccessor);
 			}
 		}
 		else {
 			simpAccessor = SimpMessageHeaderAccessor.wrap(message);
+			initHeaders(simpAccessor);
 		}
 
 		simpAccessor.setDestination(destination);
@@ -174,6 +196,11 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplate<String
 		}
 	}
 
+	private void initHeaders(SimpMessageHeaderAccessor simpAccessor) {
+		if (getHeaderInitializer() != null) {
+			getHeaderInitializer().initHeaders(simpAccessor);
+		}
+	}
 
 	@Override
 	public void convertAndSendToUser(String user, String destination, Object payload) throws MessagingException {
@@ -222,6 +249,7 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplate<String
 
 		if (headers == null) {
 			SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
+			initHeaders(headerAccessor);
 			headerAccessor.setLeaveMutable(true);
 			return headerAccessor.getMessageHeaders();
 		}
@@ -239,6 +267,7 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplate<String
 		}
 
 		SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
+		initHeaders(headerAccessor);
 		for (String key : headers.keySet()) {
 			Object value = headers.get(key);
 			headerAccessor.setNativeHeader(key, (value != null ? value.toString() : null));

spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandler.java

@@ -32,6 +32,7 @@ import org.springframework.messaging.simp.SimpMessageSendingOperations;
 import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.simp.annotation.SendToUser;
 import org.springframework.messaging.simp.user.DestinationUserNameProvider;
+import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
 
@@ -58,6 +59,8 @@ public class SendToMethodReturnValueHandler implements HandlerMethodReturnValueH
 
 	private String defaultUserDestinationPrefix = "/queue";
 
+	private MessageHeaderInitializer headerInitializer;
+
 
 	public SendToMethodReturnValueHandler(SimpMessageSendingOperations messagingTemplate, boolean annotationRequired) {
 		Assert.notNull(messagingTemplate, "messagingTemplate must not be null");
@@ -102,6 +105,23 @@ public class SendToMethodReturnValueHandler implements HandlerMethodReturnValueH
 		return this.defaultUserDestinationPrefix;
 	}
 
+	/**
+	 * Configure a {@link MessageHeaderInitializer} to apply to the headers of all
+	 * messages sent to the client outbound channel.
+	 *
+	 * <p>By default this property is not set.
+	 */
+	public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) {
+		this.headerInitializer = headerInitializer;
+	}
+
+	/**
+	 * @return the configured header initializer.
+	 */
+	public MessageHeaderInitializer getHeaderInitializer() {
+		return this.headerInitializer;
+	}
+
 
 	@Override
 	public boolean supportsReturnType(MethodParameter returnType) {
@@ -166,6 +186,9 @@ public class SendToMethodReturnValueHandler implements HandlerMethodReturnValueH
 
 	private MessageHeaders createHeaders(String sessionId) {
 		SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
+		if (getHeaderInitializer() != null) {
+			getHeaderInitializer().initHeaders(headerAccessor);
+		}
 		headerAccessor.setSessionId(sessionId);
 		headerAccessor.setLeaveMutable(true);
 		return headerAccessor.getMessageHeaders();

spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java

@@ -59,6 +59,7 @@ import org.springframework.messaging.simp.SimpMessageTypeMessageCondition;
 import org.springframework.messaging.simp.SimpMessagingTemplate;
 import org.springframework.messaging.simp.annotation.SubscribeMapping;
 import org.springframework.messaging.support.MessageHeaderAccessor;
+import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.stereotype.Controller;
 import org.springframework.util.AntPathMatcher;
 import org.springframework.util.Assert;
@@ -95,6 +96,8 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan
 
 	private Validator validator;
 
+	private MessageHeaderInitializer headerInitializer;
+
 	private final Object lifecycleMonitor = new Object();
 
 	private volatile boolean running = false;
@@ -194,6 +197,24 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan
 		this.validator = validator;
 	}
 
+	/**
+	 * Configure a {@link MessageHeaderInitializer} to pass on to
+	 * {@link org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler}s
+	 * that send messages from controller return values.
+	 *
+	 * <p>By default this property is not set.
+	 */
+	public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) {
+		this.headerInitializer = headerInitializer;
+	}
+
+	/**
+	 * @return the configured header initializer.
+	 */
+	public MessageHeaderInitializer getHeaderInitializer() {
+		return this.headerInitializer;
+	}
+
 
 	@Override
 	public boolean isAutoStartup() {
@@ -265,14 +286,21 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan
 		List<HandlerMethodReturnValueHandler> handlers = new ArrayList<HandlerMethodReturnValueHandler>();
 
 		// Annotation-based return value types
-		handlers.add(new SendToMethodReturnValueHandler(this.brokerTemplate, true));
-		handlers.add(new SubscriptionMethodReturnValueHandler(this.clientMessagingTemplate));
+		SendToMethodReturnValueHandler sth = new SendToMethodReturnValueHandler(this.brokerTemplate, true);
+		sth.setHeaderInitializer(this.headerInitializer);
+		handlers.add(sth);
+
+		SubscriptionMethodReturnValueHandler sh = new SubscriptionMethodReturnValueHandler(this.clientMessagingTemplate);
+		sh.setHeaderInitializer(this.headerInitializer);
+		handlers.add(sh);
 
 		// custom return value types
 		handlers.addAll(getCustomReturnValueHandlers());
 
 		// catch-all
-		handlers.add(new SendToMethodReturnValueHandler(this.brokerTemplate, false));
+		sth = new SendToMethodReturnValueHandler(this.brokerTemplate, false);
+		sth.setHeaderInitializer(this.headerInitializer);
+		handlers.add(sth);
 
 		return handlers;
 	}

spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java

@@ -26,6 +26,7 @@ import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.simp.annotation.SendToUser;
 import org.springframework.messaging.simp.annotation.SubscribeMapping;
+import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.util.Assert;
 
 /**
@@ -46,6 +47,8 @@ public class SubscriptionMethodReturnValueHandler implements HandlerMethodReturn
 
 	private final MessageSendingOperations<String> messagingTemplate;
 
+	private MessageHeaderInitializer headerInitializer;
+
 
 	/**
 	 * Class constructor.
@@ -58,6 +61,23 @@ public class SubscriptionMethodReturnValueHandler implements HandlerMethodReturn
 		this.messagingTemplate = messagingTemplate;
 	}
 
+	/**
+	 * Configure a {@link MessageHeaderInitializer} to apply to the headers of all
+	 * messages sent to the client outbound channel.
+	 *
+	 * <p>By default this property is not set.
+	 */
+	public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) {
+		this.headerInitializer = headerInitializer;
+	}
+
+	/**
+	 * @return the configured header initializer.
+	 */
+	public MessageHeaderInitializer getHeaderInitializer() {
+		return this.headerInitializer;
+	}
+
 
 	@Override
 	public boolean supportsReturnType(MethodParameter returnType) {
@@ -87,6 +107,9 @@ public class SubscriptionMethodReturnValueHandler implements HandlerMethodReturn
 
 	private MessageHeaders createHeaders(String sessionId, String subscriptionId) {
 		SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
+		if (getHeaderInitializer() != null) {
+			getHeaderInitializer().initHeaders(headerAccessor);
+		}
 		headerAccessor.setSessionId(sessionId);
 		headerAccessor.setSubscriptionId(subscriptionId);
 		headerAccessor.setLeaveMutable(true);

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

@@ -25,6 +25,7 @@ import org.springframework.messaging.SubscribableChannel;
 import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.util.Assert;
 import org.springframework.util.MultiValueMap;
 
@@ -48,6 +49,8 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 
 	private SubscriptionRegistry subscriptionRegistry = new DefaultSubscriptionRegistry();
 
+	private MessageHeaderInitializer headerInitializer;
+
 
 	/**
 	 * Create a SimpleBrokerMessageHandler instance with the given message channels
@@ -94,6 +97,23 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		return this.subscriptionRegistry;
 	}
 
+	/**
+	 * Configure a {@link MessageHeaderInitializer} to apply to the headers of all
+	 * messages sent to the client outbound channel.
+	 *
+	 * <p>By default this property is not set.
+	 */
+	public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) {
+		this.headerInitializer = headerInitializer;
+	}
+
+	/**
+	 * @return the configured header initializer.
+	 */
+	public MessageHeaderInitializer getHeaderInitializer() {
+		return this.headerInitializer;
+	}
+
 
 	@Override
 	public void startInternal() {
@@ -138,6 +158,7 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		}
 		else if (SimpMessageType.CONNECT.equals(messageType)) {
 			SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.create(SimpMessageType.CONNECT_ACK);
+			initHeaders(accessor);
 			accessor.setSessionId(sessionId);
 			accessor.setHeader(SimpMessageHeaderAccessor.CONNECT_MESSAGE_HEADER, message);
 			Message<byte[]> connectAck = MessageBuilder.createMessage(EMPTY_PAYLOAD, accessor.getMessageHeaders());
@@ -150,6 +171,12 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		}
 	}
 
+	private void initHeaders(SimpMessageHeaderAccessor accessor) {
+		if (getHeaderInitializer() != null) {
+			getHeaderInitializer().initHeaders(accessor);
+		}
+	}
+
 	protected void sendMessageToSubscribers(String destination, Message<?> message) {
 		MultiValueMap<String,String> subscriptions = this.subscriptionRegistry.findSubscriptions(message);
 		if ((subscriptions.size() > 0) && logger.isDebugEnabled()) {
@@ -159,6 +186,7 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		for (String sessionId : subscriptions.keySet()) {
 			for (String subscriptionId : subscriptions.get(sessionId)) {
 				SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
+				initHeaders(headerAccessor);
 				headerAccessor.setSessionId(sessionId);
 				headerAccessor.setSubscriptionId(subscriptionId);
 				headerAccessor.copyHeadersIfAbsent(message.getHeaders());

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

@@ -30,6 +30,7 @@ import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.simp.broker.AbstractBrokerMessageHandler;
 import org.springframework.messaging.support.MessageBuilder;
 import org.springframework.messaging.support.MessageHeaderAccessor;
+import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.messaging.tcp.FixedIntervalReconnectStrategy;
 import org.springframework.messaging.tcp.TcpConnection;
 import org.springframework.messaging.tcp.TcpConnectionHandler;
@@ -111,6 +112,8 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 
 	private TcpOperations<byte[]> tcpClient;
 
+	private MessageHeaderInitializer headerInitializer;
+
 	private final Map<String, StompConnectionHandler> connectionHandlers =
 			new ConcurrentHashMap<String, StompConnectionHandler>();
 
@@ -323,6 +326,24 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 		return this.tcpClient;
 	}
 
+	/**
+	 * Configure a {@link MessageHeaderInitializer} to apply to the headers of all
+	 * messages created through the {@code StompBrokerRelayMessageHandler} that
+	 * are sent to the client outbound message channel.
+	 *
+	 * <p>By default this property is not set.
+	 */
+	public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) {
+		this.headerInitializer = headerInitializer;
+	}
+
+	/**
+	 * @return the configured header initializer.
+	 */
+	public MessageHeaderInitializer getHeaderInitializer() {
+		return this.headerInitializer;
+	}
+
 
 	@Override
 	protected void startInternal() {
@@ -331,7 +352,10 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 		this.brokerChannel.subscribe(this);
 
 		if (this.tcpClient == null) {
-			this.tcpClient = new StompTcpClientFactory().create(this.relayHost, this.relayPort);
+			StompDecoder decoder = new StompDecoder();
+			decoder.setHeaderInitializer(getHeaderInitializer());
+			StompCodec codec = new StompCodec(new StompEncoder(), decoder);
+			this.tcpClient = new StompTcpClientFactory().create(this.relayHost, this.relayPort, codec);
 		}
 
 		if (logger.isDebugEnabled()) {
@@ -536,10 +560,13 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 
 		private void sendStompErrorToClient(String errorText) {
 			if (this.isRemoteClientSession) {
-				StompHeaderAccessor headers = StompHeaderAccessor.create(StompCommand.ERROR);
-				headers.setSessionId(this.sessionId);
-				headers.setMessage(errorText);
-				Message<?> errorMessage = MessageBuilder.createMessage(EMPTY_PAYLOAD, headers.getMessageHeaders());
+				StompHeaderAccessor headerAccessor = StompHeaderAccessor.create(StompCommand.ERROR);
+				if (getHeaderInitializer() != null) {
+					getHeaderInitializer().initHeaders(headerAccessor);
+				}
+				headerAccessor.setSessionId(this.sessionId);
+				headerAccessor.setMessage(errorText);
+				Message<?> errorMessage = MessageBuilder.createMessage(EMPTY_PAYLOAD, headerAccessor.getMessageHeaders());
 				sendMessageToClient(errorMessage);
 			}
 		}
@@ -811,8 +838,8 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 
 	private static class StompTcpClientFactory {
 
-		public TcpOperations<byte[]> create(String relayHost, int relayPort) {
-			return new ReactorTcpClient<byte[]>(relayHost, relayPort, new StompCodec());
+		public TcpOperations<byte[]> create(String relayHost, int relayPort, StompCodec codec) {
+			return new ReactorTcpClient<byte[]>(relayHost, relayPort, codec);
 		}
 	}
 

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

@@ -18,12 +18,13 @@ package org.springframework.messaging.simp.stomp;
 
 import org.springframework.messaging.Message;
 
+import org.springframework.util.Assert;
 import reactor.function.Consumer;
 import reactor.function.Function;
 import reactor.io.Buffer;
 import reactor.io.encoding.Codec;
 
-import java.util.List;
+import java.nio.ByteBuffer;
 
 /**
  * A Reactor TCP {@link Codec} for sending and receiving STOMP messages.
@@ -34,35 +35,68 @@ import java.util.List;
  */
 public class StompCodec implements Codec<Buffer, Message<byte[]>, Message<byte[]>> {
 
-	private static final StompDecoder DECODER = new StompDecoder();
+	private final StompDecoder stompDecoder;
 
-	private static final Function<Message<byte[]>, Buffer> ENCODER_FUNCTION = new Function<Message<byte[]>, Buffer>() {
+	private final StompEncoder stompEncoder;
 
-		private final StompEncoder encoder = new StompEncoder();
+	private final Function<Message<byte[]>, Buffer> encodingFunction;
 
-		@Override
-		public Buffer apply(Message<byte[]> message) {
-			return Buffer.wrap(this.encoder.encode(message));
-		}
-	};
+
+	public StompCodec() {
+		this(new StompEncoder(), new StompDecoder());
+	}
+
+	public StompCodec(StompEncoder encoder, StompDecoder decoder) {
+		Assert.notNull(encoder, "'encoder' is required");
+		Assert.notNull(decoder, "'decoder' is required");
+		this.stompEncoder = encoder;
+		this.stompDecoder = decoder;
+		this.encodingFunction = new EncodingFunction(this.stompEncoder);
+	}
 
 	@Override
-	public Function<Buffer, Message<byte[]>> decoder(final Consumer<Message<byte[]>> next) {
-		return new Function<Buffer, Message<byte[]>>() {
-
-			@Override
-			public Message<byte[]> apply(Buffer buffer) {
-				for (Message<byte[]> message : DECODER.decode(buffer.byteBuffer())) {
-					next.accept(message);
-				}
-				return null;
-			}
-		};
+	public Function<Buffer, Message<byte[]>> decoder(final Consumer<Message<byte[]>> messageConsumer) {
+		return new DecodingFunction(this.stompDecoder, messageConsumer);
 	}
 
 	@Override
 	public Function<Message<byte[]>, Buffer> encoder() {
-		return ENCODER_FUNCTION;
+		return this.encodingFunction;
 	}
 
+
+	private static class EncodingFunction implements Function<Message<byte[]>, Buffer> {
+
+		private final StompEncoder encoder;
+
+		private EncodingFunction(StompEncoder encoder) {
+			this.encoder = encoder;
+		}
+
+		@Override
+		public Buffer apply(Message<byte[]> message) {
+			byte[] bytes = this.encoder.encode(message);
+			return new Buffer(ByteBuffer.wrap(bytes));
+		}
+	}
+
+	private static class DecodingFunction implements Function<Buffer, Message<byte[]>> {
+
+		private final StompDecoder decoder;
+
+		private final Consumer<Message<byte[]>> messageConsumer;
+
+		public DecodingFunction(StompDecoder decoder, Consumer<Message<byte[]>> next) {
+			this.decoder = decoder;
+			this.messageConsumer = next;
+		}
+
+		@Override
+		public Message<byte[]> apply(Buffer buffer) {
+			for (Message<byte[]> message : this.decoder.decode(buffer.byteBuffer())) {
+				this.messageConsumer.accept(message);
+			}
+			return null;
+		}
+	}
 }

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

@@ -27,6 +27,7 @@ import org.apache.commons.logging.LogFactory;
 
 import org.springframework.messaging.Message;
 import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.messaging.support.NativeMessageHeaderAccessor;
 import org.springframework.util.InvalidMimeTypeException;
 import org.springframework.util.MultiValueMap;
@@ -52,6 +53,24 @@ public class StompDecoder {
 
 	private final Log logger = LogFactory.getLog(StompDecoder.class);
 
+	private MessageHeaderInitializer headerInitializer;
+
+
+	/**
+	 * Configure a
+	 * {@link org.springframework.messaging.support.MessageHeaderInitializer MessageHeaderInitializer}
+	 * to apply to the headers of {@link Message}s from decoded STOMP frames.
+	 */
+	public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) {
+		this.headerInitializer = headerInitializer;
+	}
+
+	/**
+	 * @return the configured {@code MessageHeaderInitializer} if any.
+	 */
+	public MessageHeaderInitializer getHeaderInitializer() {
+		return this.headerInitializer;
+	}
 
 	/**
 	 * Decodes one or more STOMP frames from the given {@code ByteBuffer} into a
@@ -123,6 +142,7 @@ public class StompDecoder {
 			if (buffer.remaining() > 0) {
 				StompCommand stompCommand = StompCommand.valueOf(command);
 				headerAccessor = StompHeaderAccessor.create(stompCommand);
+				initHeaders(headerAccessor);
 
 				readHeaders(buffer, headerAccessor);
 				payload = readPayload(buffer, headerAccessor);
@@ -160,12 +180,18 @@ public class StompDecoder {
 				logger.trace("Decoded heartbeat");
 			}
 			StompHeaderAccessor headerAccessor = StompHeaderAccessor.createForHeartbeat();
+			initHeaders(headerAccessor);
 			headerAccessor.setLeaveMutable(true);
 			decodedMessage = MessageBuilder.createMessage(HEARTBEAT_PAYLOAD, headerAccessor.getMessageHeaders());
 		}
 		return decodedMessage;
 	}
 
+	private void initHeaders(StompHeaderAccessor headerAccessor) {
+		if (getHeaderInitializer() != null) {
+			getHeaderInitializer().initHeaders(headerAccessor);
+		}
+	}
 
 	/**
 	 * Skip one ore more EOL characters at the start of the given ByteBuffer.

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

@@ -31,6 +31,7 @@ import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.simp.SimpMessagingTemplate;
 import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.util.Assert;
 
 /**
@@ -56,6 +57,8 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec
 
 	private final UserDestinationResolver userDestinationResolver;
 
+	private MessageHeaderInitializer headerInitializer;
+
 	private final Object lifecycleMonitor = new Object();
 
 	private volatile boolean running = false;
@@ -97,6 +100,24 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec
 		return this.userDestinationResolver;
 	}
 
+	/**
+	 * Configure a {@link MessageHeaderInitializer} to pass on to
+	 * {@link org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler}s
+	 * that send messages from controller return values.
+	 *
+	 * <p>By default this property is not set.
+	 */
+	public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) {
+		this.headerInitializer = headerInitializer;
+	}
+
+	/**
+	 * @return the configured header initializer.
+	 */
+	public MessageHeaderInitializer getHeaderInitializer() {
+		return this.headerInitializer;
+	}
+
 
 	@Override
 	public int getPhase() {
@@ -154,6 +175,9 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec
 		}
 		if (SimpMessageType.MESSAGE.equals(SimpMessageHeaderAccessor.getMessageType(message.getHeaders()))) {
 			SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.wrap(message);
+			if (getHeaderInitializer() != null) {
+				getHeaderInitializer().initHeaders(headerAccessor);
+			}
 			headerAccessor.setHeader(SimpMessageHeaderAccessor.ORIGINAL_DESTINATION, result.getSubscribeDestination());
 			message = MessageBuilder.createMessage(message.getPayload(), headerAccessor.getMessageHeaders());
 		}

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

@@ -38,66 +38,74 @@ import org.springframework.util.StringUtils;
 /**
  * A base for classes providing strongly typed getters and setters as well as
  * behavior around specific categories of headers (e.g. STOMP headers).
- * Supports creating new headers (default constructor), modifying existing headers
- * (when still mutable), or copying and modifying existing headers.
+ * Supports creating new headers, modifying existing headers (when still mutable),
+ * or copying and modifying existing headers.
  *
- * <p>The {@link #getMessageHeaders()} method provides access to the underlying,
- * fully-prepared {@code MessageHeaders} instance that can then be used as-is
- * to create a single message as follows:
+ * <p>The method {@link #getMessageHeaders()} provides access to the underlying,
+ * fully-prepared {@link MessageHeaders} that can then be used as-is (i.e.
+ * without copying) to create a single message as follows:
  *
  * <pre class="code">
- * MessageHeaderAccessor headerAccessor = new MessageHeaderAccessor();
- * headerAccessor.set("foo", "bar");
- * Message message = MessageBuilder.createMessage("payload", headerAccessor.getMessageHeaders());
+ * MessageHeaderAccessor accessor = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar");
+ * Message message = MessageBuilder.createMessage("payload", accessor.getMessageHeaders());
  * </pre>
  *
- * <p>After the above is executed, by default the {@code MessageHeaderAccessor}
- * is immutable. It is also possible to leave it mutable, for example for further
- * initialization of the message in the same thread:
+ * <p>After the above, by default the {@code MessageHeaderAccessor} becomes
+ * immutable. However it is possible to leave it mutable for further initialization
+ * in the same thread, for example:
  *
  * <pre class="code">
- * MessageHeaderAccessor headerAccessor = new MessageHeaderAccessor();
- * headerAccessor.set("foo", "bar");
- * headerAccessor.setLeaveMutable(true);
- * Message message = MessageBuilder.createMessage("payload", headerAccessor.getMessageHeaders());
+ * MessageHeaderAccessor accessor = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar");
+ * accessor.setLeaveMutable(true);
+ * Message message = MessageBuilder.createMessage("payload", accessor.getMessageHeaders());
  *
  * // later on in the same thread...
  *
- * MessageHeaderAccessor headerAccessor = MessageHeaderAccessor.getAccessor(message);
- * headerAccessor.set("bar", "baz");
- * headerAccessor.setImmutable();
+ * MessageHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message);
+ * accessor.set("bar", "baz");
+ * accessor.setImmutable();
  * </pre>
  *
- * <p>To re-obtain the {@code MessageHeaderAccessor} for a {@code MessageHeaders}
- * instance, use the following:
- *
+ * <p>The method {@link #toMap()} returns a copy of the underlying headers. It can
+ * be used to prepare multiple messages from the same {@code MessageHeaderAccessor}
+ * instance:
  * <pre class="code">
- * MessageHeaderAccessor headerAccessor = new MessageHeaderAccessor();
- * headerAccessor.set("foo", "bar");
- * Message message = MessageBuilder.createMessage("payload", headerAccessor.getMessageHeaders());
+ * MessageHeaderAccessor accessor = new MessageHeaderAccessor();
+ * MessageBuilder builder = MessageBuilder.withPayload("payload").setHeaders(accessor);
  *
- * // later on (any thread)...
- * MessageHeaderAccessor headerAccessor = MessageHeaderAccessor.getAccessor(message);
- * headerAccessor.get("foo");
- * </pre>
- *
- * <p>To prepare multiple messages with the same {@code MessageHeaderAccessor}
- * instance, use the code below. However note that this usage style does not allow
- * re-obtaining a header accessor later on:
- * <pre class="code">
- * MessageHeaderAccessor headerAccessor = new MessageHeaderAccessor();
- * MessageBuilder builder = MessageBuilder.withPayload("payload").setHeaders(headerAccessor);
- *
- * headerAccessor.setHeader("foo", "bar1");
+ * accessor.setHeader("foo", "bar1");
  * Message message1 = builder.build();
  *
- * headerAccessor.setHeader("foo", "bar2");
+ * accessor.setHeader("foo", "bar2");
  * Message message2 = builder.build();
  *
- * headerAccessor.setHeader("foo", "bar3");
+ * accessor.setHeader("foo", "bar3");
  * Message  message3 = builder.build();
  * </pre>
  *
+ * <p>However note that with the above style, the header accessor is shared and
+ * cannot be re-obtained later on. Alternatively it is also possible to create
+ * one {@code MessageHeaderAccessor} per message:
+ *
+ * <pre class="code">
+ * MessageHeaderAccessor accessor1 = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar1");
+ * Message message1 = MessageBuilder.createMessage("payload", accessor1.getMessageHeaders());
+ *
+ * MessageHeaderAccessor accessor2 = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar2");
+ * Message message2 = MessageBuilder.createMessage("payload", accessor2.getMessageHeaders());
+ *
+ * MessageHeaderAccessor accessor3 = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar3");
+ * Message message3 = MessageBuilder.createMessage("payload", accessor3.getMessageHeaders());
+ * </pre>
+ *
+ * <p>Note that the above examples aim to demonstrate the general idea of using
+ * header accessors. The most likely usage however is through sub-classes.
+ *
  * @author Rossen Stoyanchev
  * @since 4.0
  */
@@ -210,7 +218,6 @@ public class MessageHeaderAccessor {
 	 * use {@link #toMap()}.
 	 */
 	public MessageHeaders getMessageHeaders() {
-		this.headers.setIdAndTimestamp();
 		if (!this.leaveMutable) {
 			setImmutable();
 		}
@@ -259,6 +266,7 @@ public class MessageHeaderAccessor {
 	 * @since 4.1
 	 */
 	public void setImmutable() {
+		this.headers.setIdAndTimestamp();
 		this.headers.setImmutable();
 	}
 
@@ -500,6 +508,9 @@ public class MessageHeaderAccessor {
 		}
 
 		public void setIdAndTimestamp() {
+			if (!isMutable()) {
+				return;
+			}
 			if (getId() == null) {
 				IdGenerator idGenerator = (MessageHeaderAccessor.this.idGenerator != null) ?
 						MessageHeaderAccessor.this.idGenerator :