Wrapping a DestinationResolver in an aspect that adds an interceptor if it's not there already

fixes gh-838

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceSpringIntegrationAutoConfiguration.java

@@ -16,6 +16,12 @@
 
 package org.springframework.cloud.sleuth.instrument.messaging;
 
+import java.util.List;
+
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.aspectj.lang.annotation.Around;
+import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.annotation.Pointcut;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
@@ -27,7 +33,10 @@ import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.integration.channel.ChannelInterceptorAware;
 import org.springframework.integration.config.GlobalChannelInterceptor;
+import org.springframework.messaging.support.ChannelInterceptor;
+import org.springframework.messaging.support.InterceptableChannel;
 
 /**
  * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration
@@ -54,4 +63,47 @@ public class TraceSpringIntegrationAutoConfiguration {
 		return new IntegrationTraceChannelInterceptor(beanFactory);
 	}
 
+	@Bean TraceDestinationResolverAspect traceDestinationResolverAspect(TraceChannelInterceptor interceptor) {
+		return new TraceDestinationResolverAspect(interceptor);
+	}
 }
+
+@Aspect
+class TraceDestinationResolverAspect {
+
+	private final TraceChannelInterceptor interceptor;
+
+	TraceDestinationResolverAspect(TraceChannelInterceptor interceptor) {
+		this.interceptor = interceptor;
+	}
+
+	@Pointcut("execution(public * org.springframework.messaging.core.DestinationResolver.resolveDestination(..))")
+	private void anyDestinationResolver() { } // NOSONAR
+
+	@Around("anyDestinationResolver()")
+	@SuppressWarnings("unchecked")
+	public Object addInterceptor(ProceedingJoinPoint pjp) throws Throwable {
+		Object destination = pjp.proceed();
+		if (destination instanceof ChannelInterceptorAware) {
+			ChannelInterceptorAware interceptorAware = (ChannelInterceptorAware) destination;
+			if (!hasTracedChannelInterceptor(interceptorAware.getChannelInterceptors())) {
+				interceptorAware.addInterceptor(this.interceptor);
+			}
+		} else if (destination instanceof InterceptableChannel) {
+			InterceptableChannel interceptorAware = (InterceptableChannel) destination;
+			if (!hasTracedChannelInterceptor(interceptorAware.getInterceptors())) {
+				interceptorAware.addInterceptor(this.interceptor);
+			}
+		}
+		return destination;
+	}
+
+	private boolean hasTracedChannelInterceptor(List<ChannelInterceptor> interceptors) {
+		for (ChannelInterceptor channelInterceptor : interceptors) {
+			if (channelInterceptor instanceof TraceChannelInterceptor) {
+				return true;
+			}
+		}
+		return false;
+	}
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceChannelInterceptorTests.java

@@ -20,7 +20,6 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executors;
-import java.util.concurrent.TimeUnit;
 
 import org.assertj.core.api.BDDAssertions;
 import org.awaitility.Awaitility;
@@ -28,7 +27,7 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-
+import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -51,12 +50,14 @@ import org.springframework.integration.channel.ExecutorChannel;
 import org.springframework.integration.channel.QueueChannel;
 import org.springframework.integration.core.MessagingTemplate;
 import org.springframework.integration.support.MessageBuilder;
+import org.springframework.integration.support.channel.BeanFactoryChannelResolver;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.MessageHandler;
 import org.springframework.messaging.MessageHeaders;
 import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.PollableChannel;
+import org.springframework.messaging.core.DestinationResolver;
 import org.springframework.messaging.support.ChannelInterceptor;
 import org.springframework.messaging.support.ChannelInterceptorAdapter;
 import org.springframework.messaging.support.ErrorMessage;
@@ -96,6 +97,9 @@ public class TraceChannelInterceptorTests implements MessageHandler {
 	@Qualifier("tracedPollableChannel")
 	private QueueChannel queueChannel;
 
+	@Autowired
+	private DestinationResolver<MessageChannel> resolver;
+
 	@Autowired
 	private Tracer tracer;
 
@@ -189,6 +193,20 @@ public class TraceChannelInterceptorTests implements MessageHandler {
 				.setHeader(TraceMessageHeaders.SPAN_ID_NAME, Span.idToHex(20L)).build());
 		then(this.message).isNotNull();
 
+		thenMessageHasAllLogs();
+	}
+
+	@Test
+	public void addsChannelInterceptorForDynamicDestinations() {
+		this.resolver.resolveDestination("tracedChannel").send(MessageBuilder.withPayload("hi")
+				.setHeader(TraceMessageHeaders.TRACE_ID_NAME, Span.idToHex(10L))
+				.setHeader(TraceMessageHeaders.SPAN_ID_NAME, Span.idToHex(20L)).build());
+		then(this.message).isNotNull();
+
+		thenMessageHasAllLogs();
+	}
+
+	private void thenMessageHasAllLogs() {
 		String spanId = this.message.getHeaders().get(TraceMessageHeaders.SPAN_ID_NAME, String.class);
 		then(spanId).isNotNull();
 		long traceId = Span
@@ -434,27 +452,27 @@ public class TraceChannelInterceptorTests implements MessageHandler {
 		}
 
 		@Bean
-		public ExecutorChannel tracedExecutorChannel() {
+		ExecutorChannel tracedExecutorChannel() {
 			return new ExecutorChannel(Executors.newSingleThreadExecutor());
 		}
 
 		@Bean
-		public PollableChannel tracedPollableChannel() {
+		PollableChannel tracedPollableChannel() {
 			return new QueueChannel(10);
 		}
 
 		@Bean
-		public DirectChannel tracedChannel() {
+		DirectChannel tracedChannel() {
 			return new DirectChannel();
 		}
 
 		@Bean
-		public DirectChannel ignoredChannel() {
+		DirectChannel ignoredChannel() {
 			return new DirectChannel();
 		}
 
 		@Bean
-		public MessagingTemplate messagingTemplate() {
+		MessagingTemplate messagingTemplate() {
 			return new MessagingTemplate(tracedChannel());
 		}
 
@@ -463,5 +481,10 @@ public class TraceChannelInterceptorTests implements MessageHandler {
 			return new AlwaysSampler();
 		}
 
+		@Bean
+		BeanFactoryChannelResolver resolver(BeanFactory beanFactory) {
+			return new BeanFactoryChannelResolver(beanFactory);
+		}
+
 	}
 }