diff --git 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 index 5b674b0d0..38aace573 100644 --- 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 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 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(); + } + } } \ No newline at end of file diff --git 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 index 3e654945d..9f582ab94 100644 --- 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); + } } \ No newline at end of file