Added span creation for a kafka listener; fixes gh-900

2018-03-13 14:52:47 +01:00
parent 404bd9d7e4
commit 9a085e3b41
2 changed files with 134 additions and 21 deletions
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceMessagingAutoConfiguration.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceMessagingAutoConfiguration.java
@@ -16,10 +16,21 @@

 package org.springframework.cloud.sleuth.instrument.messaging;

+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.Optional;
+
+import brave.Span;
+import brave.Tracer;
 import brave.Tracing;
 import brave.kafka.clients.KafkaTracing;
 import brave.spring.rabbit.SpringRabbitTracing;
+import org.aopalliance.intercept.MethodInterceptor;
+import org.aopalliance.intercept.MethodInvocation;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.Producer;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
@@ -27,6 +38,7 @@ import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Pointcut;
 import org.springframework.amqp.rabbit.config.SimpleRabbitListenerContainerFactory;
 import org.springframework.amqp.rabbit.core.RabbitTemplate;
+import org.springframework.aop.framework.ProxyFactoryBean;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.config.BeanPostProcessor;
@@ -40,6 +52,12 @@ import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.kafka.core.ProducerFactory;
+import org.springframework.kafka.listener.AbstractMessageListenerContainer;
+import org.springframework.kafka.listener.MessageListener;
+import org.springframework.kafka.listener.MessageListenerContainer;
+import org.springframework.kafka.listener.adapter.MessagingMessageListenerAdapter;
+import org.springframework.kafka.support.converter.RecordMessageConverter;
+import org.springframework.util.ReflectionUtils;

 /**
 * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration
@@ -94,8 +112,8 @@ public class TraceMessagingAutoConfiguration {
 		@Bean
 		// for tests
 		@ConditionalOnMissingBean
-		SleuthKafkaAspect sleuthKafkaAspect(KafkaTracing kafkaTracing) {
-			return new SleuthKafkaAspect(kafkaTracing);
+		SleuthKafkaAspect sleuthKafkaAspect(KafkaTracing kafkaTracing, Tracer tracer) {
+			return new SleuthKafkaAspect(kafkaTracing, tracer);
 		}
 	}
 }
@@ -132,10 +150,16 @@ class SleuthRabbitBeanPostProcessor implements BeanPostProcessor {
@Aspect
 class SleuthKafkaAspect {

-	private final KafkaTracing kafkaTracing;
+	private static final Log log = LogFactory.getLog(SleuthKafkaAspect.class);

-	SleuthKafkaAspect(KafkaTracing kafkaTracing) {
+	private final KafkaTracing kafkaTracing;
+	private final Tracer tracer;
+	final Field recordMessageConverter;
+
+	SleuthKafkaAspect(KafkaTracing kafkaTracing, Tracer tracer) {
 		this.kafkaTracing = kafkaTracing;
+		this.tracer = tracer;
+		this.recordMessageConverter = ReflectionUtils.findField(MessagingMessageListenerAdapter.class, "recordMessageConverter");
 	}

 	@Pointcut("execution(public * org.springframework.kafka.core.ProducerFactory.createProducer(..))")
@@ -144,6 +168,9 @@ class SleuthKafkaAspect {
 	@Pointcut("execution(public * org.springframework.kafka.core.ConsumerFactory.createConsumer(..))")
 	private void anyConsumerFactory() { } // NOSONAR

+	@Pointcut("execution(public * org.springframework.kafka.config.KafkaListenerContainerFactory.createListenerContainer(..))")
+	private void anyCreateListenerContainer() { } // NOSONAR
+
 	@Around("anyProducerFactory()")
 	public Object wrapProducerFactory(ProceedingJoinPoint pjp) throws Throwable {
 		Producer producer = (Producer) pjp.proceed();
@@ -155,4 +182,85 @@ class SleuthKafkaAspect {
 		Consumer consumer = (Consumer) pjp.proceed();
 		return this.kafkaTracing.consumer(consumer);
 	}
+
+	@Around("anyCreateListenerContainer()")
+	public Object wrapListenerContainerCreation(ProceedingJoinPoint pjp) throws Throwable {
+		MessageListenerContainer listener = (MessageListenerContainer) pjp.proceed();
+		if (listener instanceof AbstractMessageListenerContainer) {
+			AbstractMessageListenerContainer container = (AbstractMessageListenerContainer) listener;
+			Object someMessageListener = container.getContainerProperties().getMessageListener();
+			if (someMessageListener == null) {
+				if (log.isDebugEnabled()) {
+					log.debug("No message listener to wrap. Proceeding");
+				}
+			} else if (someMessageListener instanceof MessageListener) {
+				container.setupMessageListener(createProxy(someMessageListener));
+			} else {
+				if (log.isDebugEnabled()) {
+					log.debug("ATM we don't support Batch message listeners");
+				}
+			}
+		} else {
+			if (log.isDebugEnabled()) {
+				log.debug("Can't wrap this listener. Proceeding");
+			}
+		}
+		return listener;
+	}
+
+	private RecordMessageConverter currentRecordMessageConverter(MessagingMessageListenerAdapter adapter)
+			throws IllegalAccessException {
+		if (this.recordMessageConverter != null) {
+			return (RecordMessageConverter) this.recordMessageConverter.get(adapter);
+		}
+		return null;
+	}
+
+	@SuppressWarnings("unchecked")
+	Object createProxy(Object bean) {
+		ProxyFactoryBean factory = new ProxyFactoryBean();
+		factory.setProxyTargetClass(true);
+		factory.addAdvice(new MessageListenerMethodInterceptor(this.kafkaTracing, this.tracer));
+		factory.setTarget(bean);
+		return factory.getObject();
+	}
+}
+
+class MessageListenerMethodInterceptor<T extends MessageListener> implements MethodInterceptor {
+
+	private static final Log log = LogFactory.getLog(MessageListenerMethodInterceptor.class);
+
+	private final KafkaTracing kafkaTracing;
+	private final Tracer tracer;
+
+	MessageListenerMethodInterceptor(KafkaTracing kafkaTracing, Tracer tracer) {
+		this.kafkaTracing = kafkaTracing;
+		this.tracer = tracer;
+	}
+
+	@Override public Object invoke(MethodInvocation invocation)
+			throws Throwable {
+		if (!"onMessage".equals(invocation.getMethod().getName())) {
+			return invocation.proceed();
+		}
+		Object[] arguments = invocation.getArguments();
+		Optional<Object> record = Arrays.stream(arguments).filter(o -> o instanceof ConsumerRecord).findFirst();
+		if (!record.isPresent()) {
+			return invocation.proceed();
+		}
+		if (log.isDebugEnabled()) {
+			log.debug("Wrapping onMessage call");
+		}
+		Span span = this.kafkaTracing.nextSpan((ConsumerRecord<?, ?>) record.get()).name("on-message").start();
+		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) {
+			return invocation.proceed();
+		} catch (RuntimeException | Error e) {
+			String message = e.getMessage();
+			if (message == null) message = e.getClass().getSimpleName();
+			span.tag("error", message);
+			throw e;
+		} finally {
+			span.finish();
+		}
+	}
 }
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceMessagingAutoConfigurationTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceMessagingAutoConfigurationTests.java
@@ -16,35 +16,26 @@

 package org.springframework.cloud.sleuth.instrument.messaging;

+import brave.Tracer;
 import brave.kafka.clients.KafkaTracing;
 import brave.sampler.Sampler;
 import brave.spring.rabbit.SpringRabbitTracing;
-import com.rabbitmq.client.Channel;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.aspectj.lang.ProceedingJoinPoint;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.BDDMockito;
-import org.mockito.Mock;
-import org.springframework.amqp.AmqpException;
-import org.springframework.amqp.core.Message;
-import org.springframework.amqp.core.MessageProperties;
-import org.springframework.amqp.rabbit.connection.CachingConnectionFactory;
-import org.springframework.amqp.rabbit.connection.Connection;
-import org.springframework.amqp.rabbit.connection.ConnectionFactory;
-import org.springframework.amqp.rabbit.connection.ConnectionListener;
 import org.springframework.amqp.rabbit.core.RabbitTemplate;
-import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
-import org.springframework.boot.test.mock.mockito.SpyBean;
 import org.springframework.cloud.sleuth.util.ArrayListSpanReporter;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.kafka.annotation.KafkaListener;
 import org.springframework.kafka.core.ConsumerFactory;
 import org.springframework.kafka.core.ProducerFactory;
+import org.springframework.kafka.listener.adapter.MessagingMessageListenerAdapter;
 import org.springframework.test.context.junit4.SpringRunner;

 import static org.assertj.core.api.BDDAssertions.then;
@@ -77,6 +68,8 @@ public class TraceMessagingAutoConfigurationTests {

 		this.consumerFactory.createConsumer();
 		then(this.mySleuthKafkaAspect.consumerWrapped).isTrue();
+
+		then(this.mySleuthKafkaAspect.adapterWrapped).isTrue();
 	}

 	@Configuration
@@ -93,8 +86,13 @@ public class TraceMessagingAutoConfigurationTests {
 		@Bean SleuthRabbitBeanPostProcessor postProcessor(BeanFactory beanFactory) {
 			return new TestSleuthRabbitBeanPostProcessor(beanFactory);
 		}
-		@Bean SleuthKafkaAspect sleuthKafkaAspect(KafkaTracing kafkaTracing) {
-			return new MySleuthKafkaAspect(kafkaTracing);
+		@Bean SleuthKafkaAspect sleuthKafkaAspect(KafkaTracing kafkaTracing, Tracer tracer) {
+			return new MySleuthKafkaAspect(kafkaTracing, tracer);
+		}
+
+		@KafkaListener(topics = "backend")
+		public void onMessage(ConsumerRecord<?, ?> message) {
+			System.err.println(message);
 		}
 	}
 }
@@ -117,9 +115,10 @@ class MySleuthKafkaAspect extends SleuthKafkaAspect {

 	boolean producerWrapped;
 	boolean consumerWrapped;
+	boolean adapterWrapped;

-	MySleuthKafkaAspect(KafkaTracing kafkaTracing) {
-		super(kafkaTracing);
+	MySleuthKafkaAspect(KafkaTracing kafkaTracing, Tracer tracer) {
+		super(kafkaTracing, tracer);
 	}

 	@Override public Object wrapProducerFactory(ProceedingJoinPoint pjp)
@@ -133,4 +132,10 @@ class MySleuthKafkaAspect extends SleuthKafkaAspect {
 		this.consumerWrapped = true;
 		return super.wrapConsumerFactory(pjp);
 	}
+
+	@Override public Object wrapListenerContainerCreation(ProceedingJoinPoint pjp)
+			throws Throwable {
+		this.adapterWrapped = true;
+		return super.wrapListenerContainerCreation(pjp);
+	}
 }