Separated the propagation setter and getter

fixes gh-1797

spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/messaging/TraceSpringMessagingAutoConfiguration.java

@@ -20,7 +20,8 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
-import org.springframework.cloud.sleuth.instrument.messaging.MessageHeaderPropagation;
+import org.springframework.cloud.sleuth.instrument.messaging.MessageHeaderPropagatorGetter;
+import org.springframework.cloud.sleuth.instrument.messaging.MessageHeaderPropagatorSetter;
 import org.springframework.cloud.sleuth.propagation.Propagator;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -35,13 +36,13 @@ class TraceSpringMessagingAutoConfiguration {
 	@Bean
 	@ConditionalOnMissingBean
 	Propagator.Setter<MessageHeaderAccessor> traceMessagePropagationSetter() {
-		return MessageHeaderPropagation.INSTANCE;
+		return new MessageHeaderPropagatorSetter();
 	}
 
 	@Bean
 	@ConditionalOnMissingBean
 	Propagator.Getter<MessageHeaderAccessor> traceMessagePropagationGetter() {
-		return MessageHeaderPropagation.INSTANCE;
+		return new MessageHeaderPropagatorGetter();
 	}
 
 }

spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/instrument/messaging/TraceWebSocketAutoConfigurationTests.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2013-2020 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.autoconfig.instrument.messaging;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.cloud.sleuth.autoconfig.TraceNoOpAutoConfiguration;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class TraceWebSocketAutoConfigurationTests {
+
+	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+			.withPropertyValues("spring.sleuth.noop.enabled=true").withConfiguration(
+					AutoConfigurations.of(TraceNoOpAutoConfiguration.class, TraceSpringMessagingAutoConfiguration.class,
+							TraceSpringIntegrationAutoConfiguration.class, TraceWebSocketAutoConfiguration.class));
+
+	@Test
+	void should_inject_beans_for_getter_setter_messaging_propagation() {
+		this.contextRunner.run(context -> assertThat(context).hasSingleBean(TraceWebSocketAutoConfiguration.class));
+	}
+
+}

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagatorGetter.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2013-2020 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.instrument.messaging;
+
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.springframework.cloud.sleuth.propagation.Propagator;
+import org.springframework.messaging.support.MessageHeaderAccessor;
+import org.springframework.messaging.support.NativeMessageHeaderAccessor;
+import org.springframework.util.StringUtils;
+
+/**
+ * Getter for Spring Integration based communication.
+ *
+ * This always sets native headers in defence of STOMP issues discussed <a href=
+ * "https://github.com/spring-cloud/spring-cloud-sleuth/issues/716#issuecomment-337523705">here</a>.
+ *
+ * @author Marcin Grzejszczak
+ * @since 3.0.0
+ */
+public class MessageHeaderPropagatorGetter implements Propagator.Getter<MessageHeaderAccessor> {
+
+	private static final Log log = LogFactory.getLog(MessageHeaderPropagatorGetter.class);
+
+	@Override
+	public String get(MessageHeaderAccessor accessor, String key) {
+		try {
+			String value = doGet(accessor, key);
+			if (StringUtils.hasText(value)) {
+				return value;
+			}
+		}
+		catch (Exception ex) {
+			if (log.isDebugEnabled()) {
+				log.debug("An exception happened when we tried to retrieve the [" + key + "] from message", ex);
+			}
+		}
+		return null;
+	}
+
+	private String doGet(MessageHeaderAccessor accessor, String key) {
+		if (accessor instanceof NativeMessageHeaderAccessor) {
+			NativeMessageHeaderAccessor nativeAccessor = (NativeMessageHeaderAccessor) accessor;
+			String result = nativeAccessor.getFirstNativeHeader(key);
+			if (result != null) {
+				return result;
+			}
+		}
+		else {
+			Object nativeHeaders = accessor.getHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS);
+			if (nativeHeaders instanceof Map) {
+				Object result = ((Map) nativeHeaders).get(key);
+				if (result instanceof List && !((List) result).isEmpty()) {
+					return String.valueOf(((List) result).get(0));
+				}
+			}
+		}
+		Object result = accessor.getHeader(key);
+		if (result != null) {
+			if (result instanceof byte[]) {
+				return new String((byte[]) result, StandardCharsets.UTF_8);
+			}
+			return result.toString();
+		}
+		return null;
+	}
+
+	@Override
+	public String toString() {
+		return "MessageHeaderPropagatorGetter{}";
+	}
+
+}

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessageHeaderPropagatorSetter.java

@@ -16,7 +16,6 @@
 
 package org.springframework.cloud.sleuth.instrument.messaging;
 
-import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -29,10 +28,9 @@ import org.springframework.cloud.sleuth.propagation.Propagator;
 import org.springframework.messaging.support.MessageHeaderAccessor;
 import org.springframework.messaging.support.NativeMessageHeaderAccessor;
 import org.springframework.util.LinkedMultiValueMap;
-import org.springframework.util.StringUtils;
 
 /**
- * Setter and getter for Spring Integration based communication.
+ * Setter for Spring Integration based communication.
  *
  * This always sets native headers in defence of STOMP issues discussed <a href=
  * "https://github.com/spring-cloud/spring-cloud-sleuth/issues/716#issuecomment-337523705">here</a>.
@@ -40,15 +38,9 @@ import org.springframework.util.StringUtils;
  * @author Marcin Grzejszczak
  * @since 3.0.0
  */
-public enum MessageHeaderPropagation
-		implements Propagator.Setter<MessageHeaderAccessor>, Propagator.Getter<MessageHeaderAccessor> {
+public class MessageHeaderPropagatorSetter implements Propagator.Setter<MessageHeaderAccessor> {
 
-	/**
-	 * Singleton instance for message header propagation.
-	 */
-	INSTANCE;
-
-	private static final Log log = LogFactory.getLog(MessageHeaderPropagation.class);
+	private static final Log log = LogFactory.getLog(MessageHeaderPropagatorSetter.class);
 
 	static Map<String, ?> propagationHeaders(Map<String, ?> headers, List<String> propagationHeaders) {
 		Map<String, Object> headersToCopy = new HashMap<>();
@@ -134,52 +126,9 @@ public enum MessageHeaderPropagation
 		return (map != null ? new LinkedMultiValueMap<>(map) : Collections.emptyMap());
 	}
 
-	@Override
-	public String get(MessageHeaderAccessor accessor, String key) {
-		try {
-			String value = doGet(accessor, key);
-			if (StringUtils.hasText(value)) {
-				return value;
-			}
-		}
-		catch (Exception ex) {
-			if (log.isDebugEnabled()) {
-				log.debug("An exception happened when we tried to retrieve the [" + key + "] from message", ex);
-			}
-		}
-		return null;
-	}
-
-	private String doGet(MessageHeaderAccessor accessor, String key) {
-		if (accessor instanceof NativeMessageHeaderAccessor) {
-			NativeMessageHeaderAccessor nativeAccessor = (NativeMessageHeaderAccessor) accessor;
-			String result = nativeAccessor.getFirstNativeHeader(key);
-			if (result != null) {
-				return result;
-			}
-		}
-		else {
-			Object nativeHeaders = accessor.getHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS);
-			if (nativeHeaders instanceof Map) {
-				Object result = ((Map) nativeHeaders).get(key);
-				if (result instanceof List && !((List) result).isEmpty()) {
-					return String.valueOf(((List) result).get(0));
-				}
-			}
-		}
-		Object result = accessor.getHeader(key);
-		if (result != null) {
-			if (result instanceof byte[]) {
-				return new String((byte[]) result, StandardCharsets.UTF_8);
-			}
-			return result.toString();
-		}
-		return null;
-	}
-
 	@Override
 	public String toString() {
-		return "MessageHeaderPropagation{}";
+		return "MessageHeaderPropagatorSetter{}";
 	}
 
 }

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceMessageHandler.java

@@ -268,7 +268,7 @@ class TraceMessageHandler {
 		clearTechnicalTracingHeaders(headers);
 		if (originalMessage instanceof ErrorMessage) {
 			ErrorMessage errorMessage = (ErrorMessage) originalMessage;
-			headers.copyHeaders(MessageHeaderPropagation.propagationHeaders(additionalHeaders.getMessageHeaders(),
+			headers.copyHeaders(MessageHeaderPropagatorSetter.propagationHeaders(additionalHeaders.getMessageHeaders(),
 					this.propagator.fields()));
 			return new ErrorMessage(errorMessage.getPayload(), isWebSockets(headers) ? headers.getMessageHeaders()
 					: new MessageHeaders(headers.getMessageHeaders()), errorMessage.getOriginalMessage());
@@ -307,11 +307,11 @@ class TraceMessageHandler {
 		List<String> keysToRemove = new ArrayList<>(this.propagator.fields());
 		keysToRemove.add(Span.class.getName());
 		keysToRemove.add("traceHandlerParentSpan");
-		MessageHeaderPropagation.removeAnyTraceHeaders(headers, keysToRemove);
+		MessageHeaderPropagatorSetter.removeAnyTraceHeaders(headers, keysToRemove);
 	}
 
 	private void clearTechnicalTracingHeaders(MessageHeaderAccessor headers) {
-		MessageHeaderPropagation.removeAnyTraceHeaders(headers,
+		MessageHeaderPropagatorSetter.removeAnyTraceHeaders(headers,
 				Arrays.asList(Span.class.getName(), "traceHandlerParentSpan"));
 	}
 

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java

@@ -142,7 +142,7 @@ public final class TracingChannelInterceptor extends ChannelInterceptorAdapter
 		}
 		MessageHeaderAccessor headers = mutableHeaderAccessor(retrievedMessage);
 		Span.Builder spanBuilder = this.propagator.extract(headers, this.extractor);
-		MessageHeaderPropagation.removeAnyTraceHeaders(headers, this.propagator.fields());
+		MessageHeaderPropagatorSetter.removeAnyTraceHeaders(headers, this.propagator.fields());
 		spanBuilder = spanBuilder.kind(Span.Kind.PRODUCER);
 		spanBuilder = this.messageSpanCustomizer.customizeSend(spanBuilder, message, channel)
 				.remoteServiceName(toRemoteServiceName(headers));
@@ -196,7 +196,7 @@ public final class TracingChannelInterceptor extends ChannelInterceptorAdapter
 		MessageHeaderAccessor headers = mutableHeaderAccessor(originalMessage);
 		if (originalMessage instanceof ErrorMessage) {
 			ErrorMessage errorMessage = (ErrorMessage) originalMessage;
-			headers.copyHeaders(MessageHeaderPropagation.propagationHeaders(additionalHeaders.getMessageHeaders(),
+			headers.copyHeaders(MessageHeaderPropagatorSetter.propagationHeaders(additionalHeaders.getMessageHeaders(),
 					this.propagator.fields()));
 			return new ErrorMessage(errorMessage.getPayload(), isWebSockets(headers) ? headers.getMessageHeaders()
 					: new MessageHeaders(headers.getMessageHeaders()), errorMessage.getOriginalMessage());
@@ -279,7 +279,7 @@ public final class TracingChannelInterceptor extends ChannelInterceptorAdapter
 	private Span consumerSpanReceive(Message<?> message, MessageChannel channel, MessageHeaderAccessor headers,
 			Span result) {
 		Span.Builder builder = this.tracer.spanBuilder().setParent(result.context());
-		MessageHeaderPropagation.removeAnyTraceHeaders(headers, this.propagator.fields());
+		MessageHeaderPropagatorSetter.removeAnyTraceHeaders(headers, this.propagator.fields());
 		builder = builder.kind(Span.Kind.CONSUMER);
 		builder = this.messageSpanCustomizer.customizeReceive(builder, message, channel);
 		builder = builder.remoteServiceName(toRemoteServiceName(headers));
@@ -321,7 +321,7 @@ public final class TracingChannelInterceptor extends ChannelInterceptorAdapter
 		// remove any trace headers, but don't re-inject as we are synchronously
 		// processing the
 		// message and can rely on scoping to access this span later.
-		MessageHeaderPropagation.removeAnyTraceHeaders(headers, this.propagator.fields());
+		MessageHeaderPropagatorSetter.removeAnyTraceHeaders(headers, this.propagator.fields());
 		if (log.isDebugEnabled()) {
 			log.debug("Created a new span in before handle " + handle);
 		}

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptorTest.java

@@ -74,8 +74,9 @@ public abstract class TracingChannelInterceptorTest implements TestTracingAwareS
 	}
 
 	protected ChannelInterceptor interceptor = new TracingChannelInterceptor(tracerTest().tracing().tracer(),
-			tracerTest().tracing().propagator(), MessageHeaderPropagation.INSTANCE, MessageHeaderPropagation.INSTANCE,
-			remoteServiceNameMapper(new SleuthMessagingProperties()), new DefaultMessageSpanCustomizer());
+			tracerTest().tracing().propagator(), new MessageHeaderPropagatorSetter(),
+			new MessageHeaderPropagatorGetter(), remoteServiceNameMapper(new SleuthMessagingProperties()),
+			new DefaultMessageSpanCustomizer());
 
 	protected TestSpanHandler spans = tracerTest().handler();
 
@@ -129,8 +130,8 @@ public abstract class TracingChannelInterceptorTest implements TestTracingAwareS
 	@Test
 	public void allowsSpanCustomization() {
 		this.interceptor = new TracingChannelInterceptor(tracerTest().tracing().tracer(),
-				tracerTest().tracing().propagator(), MessageHeaderPropagation.INSTANCE,
-				MessageHeaderPropagation.INSTANCE, remoteServiceNameMapper(new SleuthMessagingProperties()),
+				tracerTest().tracing().propagator(), new MessageHeaderPropagatorSetter(),
+				new MessageHeaderPropagatorGetter(), remoteServiceNameMapper(new SleuthMessagingProperties()),
 				new MyMessageSpanCustomizer());
 
 		this.directChannel.addInterceptor(this.interceptor);