Don't renew message until newMessageNextTrack()

artembilan · artembilan · commit a90d66fedced · 2016-04-07T22:14:12.000-04:00
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/channel/AbstractMessageChannel.java b/spring-integration-core/src/main/java/org/springframework/integration/channel/AbstractMessageChannel.java
@@ -98,6 +98,11 @@ public void setShouldTrack(boolean shouldTrack) {
 		this.shouldTrack = shouldTrack;
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return false;
+	}
+
 	@Override
 	public void setCountsEnabled(boolean countsEnabled) {
 		this.countsEnabled = countsEnabled;
@@ -392,7 +397,7 @@ public final boolean send(Message<?> message, long timeout) {
 		Assert.notNull(message, "message must not be null");
 		Assert.notNull(message.getPayload(), "message payload must not be null");
 		if (this.shouldTrack) {
-			message = MessageHistory.write(message, this, this.getMessageBuilderFactory());
+			message = MessageHistory.write(message, this, getMessageBuilderFactory(), newMessageNextTrack());
 		}
 
 		Deque<ChannelInterceptor> interceptorStack = null;
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/channel/PublishSubscribeChannel.java b/spring-integration-core/src/main/java/org/springframework/integration/channel/PublishSubscribeChannel.java
@@ -70,6 +70,11 @@ public String getComponentType() {
 		return "publish-subscribe-channel";
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return true;
+	}
+
 	/**
 	 * Provide an {@link ErrorHandler} strategy for handling Exceptions that
 	 * occur downstream from this channel. This will <i>only</i> be applied if
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/endpoint/MessageProducerSupport.java b/spring-integration-core/src/main/java/org/springframework/integration/endpoint/MessageProducerSupport.java
@@ -80,6 +80,11 @@ public void setShouldTrack(boolean shouldTrack) {
 		this.shouldTrack = shouldTrack;
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return false;
+	}
+
 	protected MessagingTemplate getMessagingTemplate() {
 		return this.messagingTemplate;
 	}
@@ -117,7 +122,7 @@ protected void sendMessage(Message<?> message) {
 			throw new MessagingException("cannot send a null message");
 		}
 		if (this.shouldTrack) {
-			message = MessageHistory.write(message, this, this.getMessageBuilderFactory());
+			message = MessageHistory.write(message, this, getMessageBuilderFactory(), newMessageNextTrack());
 		}
 		try {
 			this.messagingTemplate.send(this.outputChannel, message);
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/endpoint/SourcePollingChannelAdapter.java b/spring-integration-core/src/main/java/org/springframework/integration/endpoint/SourcePollingChannelAdapter.java
@@ -102,6 +102,11 @@ public void setShouldTrack(boolean shouldTrack) {
 		this.shouldTrack = shouldTrack;
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return false;
+	}
+
 	@Override
 	public String getComponentType() {
 		return (this.source instanceof NamedComponent) ?
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/gateway/GatewayProxyFactoryBean.java b/spring-integration-core/src/main/java/org/springframework/integration/gateway/GatewayProxyFactoryBean.java
@@ -219,6 +219,11 @@ public void setShouldTrack(boolean shouldTrack) {
 		}
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return false;
+	}
+
 	/**
 	 * Set the executor for use when the gateway method returns
 	 * {@link java.util.concurrent.Future} or {@link org.springframework.util.concurrent.ListenableFuture}.
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/gateway/MessagingGatewaySupport.java b/spring-integration-core/src/main/java/org/springframework/integration/gateway/MessagingGatewaySupport.java
@@ -234,6 +234,12 @@ public void setShouldTrack(boolean shouldTrack) {
 		this.historyWritingPostProcessor.setShouldTrack(shouldTrack);
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return false;
+	}
+
+
 	@Override
 	public int getMessageCount() {
 		return (int) this.messageCount.get();
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/handler/AbstractMessageHandler.java b/spring-integration-core/src/main/java/org/springframework/integration/handler/AbstractMessageHandler.java
@@ -94,6 +94,11 @@ public void setShouldTrack(boolean shouldTrack) {
 		this.shouldTrack = shouldTrack;
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return false;
+	}
+
 	@Override
 	public void configureMetrics(AbstractMessageHandlerMetrics metrics) {
 		Assert.notNull(metrics, "'metrics' must not be null");
@@ -119,7 +124,7 @@ public final void handleMessage(Message<?> message) {
 		AbstractMessageHandlerMetrics handlerMetrics = this.handlerMetrics;
 		try {
 			if (message != null && this.shouldTrack) {
-				message = MessageHistory.write(message, this, this.getMessageBuilderFactory());
+				message = MessageHistory.write(message, this, getMessageBuilderFactory(), newMessageNextTrack());
 			}
 			if (countsEnabled) {
 				start = handlerMetrics.beforeHandle();
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/history/MessageHistory.java b/spring-integration-core/src/main/java/org/springframework/integration/history/MessageHistory.java
@@ -55,7 +55,9 @@ public final class MessageHistory implements List<Properties>, Serializable {
 	private static final MessageBuilderFactory MESSAGE_BUILDER_FACTORY = new DefaultMessageBuilderFactory();
 
 
-	private final List<Properties> components;
+	private final List<Properties> components = new ArrayList<Properties>();
+
+	private boolean nextNewMessage;
 
 
 	public static MessageHistory read(Message<?> message) {
@@ -66,46 +68,58 @@ public static <T> Message<T> write(Message<T> message, NamedComponent component)
 		return write(message, component, MESSAGE_BUILDER_FACTORY);
 	}
 
-	@SuppressWarnings("unchecked")
 	public static <T> Message<T> write(Message<T> message, NamedComponent component,
 			MessageBuilderFactory messageBuilderFactory) {
+		return write(message, component, messageBuilderFactory, false);
+	}
+
+	@SuppressWarnings("unchecked")
+	public static <T> Message<T> write(Message<T> message, NamedComponent component,
+			MessageBuilderFactory messageBuilderFactory, boolean newMessageOnNext) {
 		Assert.notNull(message, "Message must not be null");
 		Assert.notNull(component, "Component must not be null");
 		Properties metadata = extractMetadata(component);
 		if (!metadata.isEmpty()) {
-			MessageHistory previousHistory = message.getHeaders().get(HEADER_NAME, MessageHistory.class);
-			List<Properties> components = (previousHistory != null) ?
-					new ArrayList<Properties>(previousHistory) : new ArrayList<Properties>();
-			components.add(metadata);
-			MessageHistory history = new MessageHistory(components);
-
-			if (message instanceof MutableMessage) {
-				message.getHeaders().put(HEADER_NAME, history);
-			}
-			else if (message instanceof ErrorMessage) {
-				IntegrationMessageHeaderAccessor headerAccessor = new IntegrationMessageHeaderAccessor(message);
-				headerAccessor.setHeader(HEADER_NAME, history);
-				Throwable payload = ((ErrorMessage) message).getPayload();
-				ErrorMessage errorMessage = new ErrorMessage(payload, headerAccessor.toMessageHeaders());
-				message = (Message<T>) errorMessage;
-			}
-			else if (message instanceof AdviceMessage) {
-				IntegrationMessageHeaderAccessor headerAccessor = new IntegrationMessageHeaderAccessor(message);
-				headerAccessor.setHeader(HEADER_NAME, history);
-				message = new AdviceMessage<T>(message.getPayload(), headerAccessor.toMessageHeaders(),
-						((AdviceMessage) message).getInputMessage());
+			MessageHistory messageHistory = read(message);
+			if (messageHistory == null || messageHistory.nextNewMessage) {
+				messageHistory = new MessageHistory(messageHistory);
+				messageHistory.components.add(metadata);
+
+				if (message instanceof MutableMessage) {
+					message.getHeaders().put(HEADER_NAME, messageHistory);
+				}
+				else if (message instanceof ErrorMessage) {
+					IntegrationMessageHeaderAccessor headerAccessor = new IntegrationMessageHeaderAccessor(message);
+					headerAccessor.setHeader(HEADER_NAME, messageHistory);
+					Throwable payload = ((ErrorMessage) message).getPayload();
+					ErrorMessage errorMessage = new ErrorMessage(payload, headerAccessor.toMessageHeaders());
+					message = (Message<T>) errorMessage;
+				}
+				else if (message instanceof AdviceMessage) {
+					IntegrationMessageHeaderAccessor headerAccessor = new IntegrationMessageHeaderAccessor(message);
+					headerAccessor.setHeader(HEADER_NAME, messageHistory);
+					message = new AdviceMessage<T>(message.getPayload(), headerAccessor.toMessageHeaders(),
+							((AdviceMessage) message).getInputMessage());
+				}
+				else {
+					message = messageBuilderFactory.fromMessage(message)
+							.setHeader(HEADER_NAME, messageHistory)
+							.build();
+				}
 			}
 			else {
-				message = messageBuilderFactory.fromMessage(message).setHeader(HEADER_NAME, history).build();
+				messageHistory.components.add(metadata);
 			}
+			messageHistory.nextNewMessage = newMessageOnNext;
 		}
 		return message;
 	}
 
 
 	private MessageHistory(List<Properties> components) {
-		Assert.notEmpty(components, "component list must not be empty");
-		this.components = components;
+		if (components != null) {
+			this.components.addAll(components);
+		}
 	}
 
 
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/router/AbstractMessageRouter.java b/spring-integration-core/src/main/java/org/springframework/integration/router/AbstractMessageRouter.java
@@ -121,6 +121,11 @@ public String getComponentType() {
 		return "router";
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return true;
+	}
+
 	/**
 	 * Provides {@link MessagingTemplate} access for subclasses
 	 * @return The messaging template.
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/support/management/LifecycleTrackableMessageHandlerMetrics.java b/spring-integration-core/src/main/java/org/springframework/integration/support/management/LifecycleTrackableMessageHandlerMetrics.java
@@ -52,4 +52,9 @@ public void setShouldTrack(boolean shouldTrack) {
 		this.trackable.setShouldTrack(shouldTrack);
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return this.trackable.newMessageNextTrack();
+	}
+
 }
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/support/management/LifecycleTrackableMessageSourceMetrics.java b/spring-integration-core/src/main/java/org/springframework/integration/support/management/LifecycleTrackableMessageSourceMetrics.java
@@ -52,4 +52,9 @@ public void setShouldTrack(boolean shouldTrack) {
 		this.trackable.setShouldTrack(shouldTrack);
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return this.trackable.newMessageNextTrack();
+	}
+
 }
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/support/management/TrackableComponent.java b/spring-integration-core/src/main/java/org/springframework/integration/support/management/TrackableComponent.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2015 the original author or authors.
+ * Copyright 2002-2016 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,6 +22,7 @@
 /**
  * @author Mark Fisher
  * @author Gary Russell
+ * @author Artem Bilan
  * @since 2.0
  */
 @IntegrationManagedResource
@@ -30,4 +31,6 @@ public interface TrackableComponent extends NamedComponent {
 	@ManagedOperation
 	void setShouldTrack(boolean shouldTrack);
 
+	boolean newMessageNextTrack();
+
 }
diff --git a/spring-integration-core/src/main/java/org/springframework/integration/support/management/TrackableRouterMetrics.java b/spring-integration-core/src/main/java/org/springframework/integration/support/management/TrackableRouterMetrics.java
@@ -50,4 +50,9 @@ public void setShouldTrack(boolean shouldTrack) {
 		this.trackable.setShouldTrack(shouldTrack);
 	}
 
+	@Override
+	public boolean newMessageNextTrack() {
+		return this.trackable.newMessageNextTrack();
+	}
+
 }
diff --git a/spring-integration-core/src/test/java/org/springframework/integration/core/MessageHistoryTests.java b/spring-integration-core/src/test/java/org/springframework/integration/core/MessageHistoryTests.java
@@ -25,9 +25,12 @@
 import static org.junit.Assert.assertThat;
 
 import java.util.Properties;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.Test;
 
+import org.springframework.integration.channel.PublishSubscribeChannel;
+import org.springframework.integration.handler.AbstractMessageHandler;
 import org.springframework.integration.history.MessageHistory;
 import org.springframework.integration.message.AdviceMessage;
 import org.springframework.integration.support.MessageBuilder;
@@ -53,6 +56,7 @@ public void addComponents() {
 		assertNotNull(history1);
 		assertEquals("testComponent-1", history1.toString());
 		Message<String> result2 = MessageHistory.write(result1, new TestComponent(2));
+		assertSame(result1, result2);
 		MessageHistory history2 = MessageHistory.read(result2);
 		assertNotNull(history2);
 		assertEquals("testComponent-1,testComponent-2", history2.toString());
@@ -97,7 +101,7 @@ public void testCorrectErrorMessageAfterWrite() {
 		Message<Throwable> result2 = MessageHistory.write(result1, new TestComponent(2));
 		assertThat(result2, instanceOf(ErrorMessage.class));
 		assertNotSame(original, result2);
-		assertNotSame(result1, result2);
+		assertSame(result1, result2);
 		assertSame(original.getPayload(), result2.getPayload());
 		MessageHistory history2 = MessageHistory.read(result2);
 		assertNotNull(history2);
@@ -122,12 +126,63 @@ public void testCorrectAdviceMessageAfterWrite() {
 		assertNotSame(original, result2);
 		assertSame(original.getPayload(), result2.getPayload());
 		assertSame(original.getInputMessage(), ((AdviceMessage) result2).getInputMessage());
-		assertNotSame(result1, result2);
+		assertSame(result1, result2);
 		MessageHistory history2 = MessageHistory.read(result2);
 		assertNotNull(history2);
 		assertEquals("testComponent-1,testComponent-2", history2.toString());
 	}
 
+	@Test
+	public void testHistoryOnPublishSubscribeChannel() {
+		Message<?> message = new GenericMessage<>("foo");
+		assertNull(MessageHistory.read(message));
+
+		PublishSubscribeChannel publishSubscribeChannel = new PublishSubscribeChannel();
+		publishSubscribeChannel.setComponentName("publishSubscribeChannel");
+		publishSubscribeChannel.setShouldTrack(true);
+
+		final AtomicReference<Message<?>> message1 = new AtomicReference<>();
+
+		AbstractMessageHandler messageHandler1 = new AbstractMessageHandler() {
+
+			@Override
+			protected void handleMessageInternal(Message<?> message) throws Exception {
+				message1.set(message);
+			}
+
+		};
+
+		messageHandler1.setShouldTrack(true);
+		messageHandler1.setComponentName("messageHandler1");
+		publishSubscribeChannel.subscribe(messageHandler1);
+
+		final AtomicReference<Message<?>> message2 = new AtomicReference<>();
+
+		AbstractMessageHandler messageHandler2 = new AbstractMessageHandler() {
+
+			@Override
+			protected void handleMessageInternal(Message<?> message) throws Exception {
+				message2.set(message);
+			}
+
+		};
+
+		messageHandler2.setShouldTrack(true);
+		messageHandler2.setComponentName("messageHandler2");
+		publishSubscribeChannel.subscribe(messageHandler2);
+
+		publishSubscribeChannel.send(message);
+
+		MessageHistory history1 = MessageHistory.read(message1.get());
+		assertNotNull(history1);
+		assertEquals("publishSubscribeChannel,messageHandler1", history1.toString());
+
+		MessageHistory history2 = MessageHistory.read(message2.get());
+		assertNotNull(history2);
+		assertEquals("publishSubscribeChannel,messageHandler2", history2.toString());
+	}
+
+
 
 	private static class TestComponent implements NamedComponent {
 
diff --git a/spring-integration-core/src/test/java/org/springframework/integration/router/RecipientListRouterTests.java b/spring-integration-core/src/test/java/org/springframework/integration/router/RecipientListRouterTests.java

Original file line number	Diff line number	Diff line change
`@@ -80,6 +80,11 @@ public void setShouldTrack(boolean shouldTrack) {`
`80`	`80`	`this.shouldTrack = shouldTrack;`
`81`	`81`	`}`
`82`	`82`
	`83`	`+ @Override`
	`84`	`+ public boolean newMessageNextTrack() {`
	`85`	`+ return false;`
	`86`	`+ }`
	`87`	`+`
`83`	`88`	`protected MessagingTemplate getMessagingTemplate() {`
`84`	`89`	`return this.messagingTemplate;`
`85`	`90`	`}`
`@@ -117,7 +122,7 @@ protected void sendMessage(Message<?> message) {`
`117`	`122`	`throw new MessagingException("cannot send a null message");`
`118`	`123`	`}`
`119`	`124`	`if (this.shouldTrack) {`
`120`		`- message = MessageHistory.write(message, this, this.getMessageBuilderFactory());`
	`125`	`+ message = MessageHistory.write(message, this, getMessageBuilderFactory(), newMessageNextTrack());`
`121`	`126`	`}`
`122`	`127`	`try {`
`123`	`128`	`this.messagingTemplate.send(this.outputChannel, message);`
Original file line number	Diff line number	Diff line change
`@@ -219,6 +219,11 @@ public void setShouldTrack(boolean shouldTrack) {`
`219`	`219`	`}`
`220`	`220`	`}`
`221`	`221`
	`222`	`+ @Override`
	`223`	`+ public boolean newMessageNextTrack() {`
	`224`	`+ return false;`
	`225`	`+ }`
	`226`	`+`
`222`	`227`	`/**`
`223`	`228`	`* Set the executor for use when the gateway method returns`
`224`	`229`	`* {@link java.util.concurrent.Future} or {@link org.springframework.util.concurrent.ListenableFuture}.`
Original file line number	Diff line number	Diff line change
`@@ -52,4 +52,9 @@ public void setShouldTrack(boolean shouldTrack) {`
`52`	`52`	`this.trackable.setShouldTrack(shouldTrack);`
`53`	`53`	`}`
`54`	`54`
	`55`	`+ @Override`
	`56`	`+ public boolean newMessageNextTrack() {`
	`57`	`+ return this.trackable.newMessageNextTrack();`
	`58`	`+ }`
	`59`	`+`
`55`	`60`	`}`