From 0b5dec8cd7bc6864be1510bc04910f9bc8a7dd82 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Fri, 31 Jul 2020 09:57:03 +0200 Subject: [PATCH] Added Consumer / Producer post processors (#1701) fixes gh-1670 --- .../TraceMessagingAutoConfiguration.java | 90 +++++++++++++ .../KafkaFactoryBeanPostProcessorTests.java | 127 ++++++++++++++++++ 2 files changed, 217 insertions(+) create mode 100644 spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/KafkaFactoryBeanPostProcessorTests.java 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 a8fa498ab..eec48bd12 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 @@ -61,7 +61,10 @@ import org.springframework.context.annotation.Role; import org.springframework.jms.annotation.JmsListenerConfigurer; import org.springframework.jms.config.JmsListenerEndpointRegistry; import org.springframework.jms.config.TracingJmsListenerEndpointRegistry; +import org.springframework.kafka.core.ConsumerFactory; +import org.springframework.kafka.core.ConsumerPostProcessor; import org.springframework.kafka.core.ProducerFactory; +import org.springframework.kafka.core.ProducerPostProcessor; import org.springframework.kafka.listener.AbstractMessageListenerContainer; import org.springframework.kafka.listener.MessageListener; import org.springframework.kafka.listener.MessageListenerContainer; @@ -159,6 +162,12 @@ public class TraceMessagingAutoConfiguration { return new SleuthKafkaAspect(kafkaTracing, tracer); } + @Bean + KafkaFactoryBeanPostProcessor kafkaFactoryBeanPostProcessor( + BeanFactory beanFactory) { + return new KafkaFactoryBeanPostProcessor(beanFactory); + } + } @Configuration(proxyBeanMethods = false) @@ -241,6 +250,87 @@ class SleuthRabbitBeanPostProcessor implements BeanPostProcessor { } +class KafkaFactoryBeanPostProcessor implements BeanPostProcessor { + + private final BeanFactory beanFactory; + + KafkaFactoryBeanPostProcessor(BeanFactory beanFactory) { + this.beanFactory = beanFactory; + } + + @SuppressWarnings("unchecked") + @Override + public Object postProcessAfterInitialization(Object bean, String beanName) + throws BeansException { + if (bean instanceof ConsumerFactory) { + ConsumerFactory factory = (ConsumerFactory) bean; + if (factory.getPostProcessors().stream() + .noneMatch(o -> o instanceof TraceConsumerPostProcessor)) { + factory.addPostProcessor( + new TraceConsumerPostProcessor(this.beanFactory)); + } + } + else if (bean instanceof ProducerFactory) { + ProducerFactory factory = (ProducerFactory) bean; + if (factory.getPostProcessors().stream() + .noneMatch(o -> o instanceof TraceProducerPostProcessor)) { + factory.addPostProcessor( + new TraceProducerPostProcessor(this.beanFactory)); + } + } + return bean; + } + +} + +class TraceConsumerPostProcessor implements ConsumerPostProcessor { + + private final BeanFactory beanFactory; + + private KafkaTracing kafkaTracing; + + TraceConsumerPostProcessor(BeanFactory beanFactory) { + this.beanFactory = beanFactory; + } + + private KafkaTracing kafkaTracing() { + if (this.kafkaTracing == null) { + this.kafkaTracing = this.beanFactory.getBean(KafkaTracing.class); + } + return this.kafkaTracing; + } + + @Override + public Consumer apply(Consumer kvConsumer) { + return kafkaTracing().consumer(kvConsumer); + } + +} + +class TraceProducerPostProcessor implements ProducerPostProcessor { + + private final BeanFactory beanFactory; + + private KafkaTracing kafkaTracing; + + TraceProducerPostProcessor(BeanFactory beanFactory) { + this.beanFactory = beanFactory; + } + + private KafkaTracing kafkaTracing() { + if (this.kafkaTracing == null) { + this.kafkaTracing = this.beanFactory.getBean(KafkaTracing.class); + } + return this.kafkaTracing; + } + + @Override + public Producer apply(Producer kvProducer) { + return kafkaTracing().producer(kvProducer); + } + +} + @Aspect class SleuthKafkaAspect { diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/KafkaFactoryBeanPostProcessorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/KafkaFactoryBeanPostProcessorTests.java new file mode 100644 index 000000000..670ed02f0 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/KafkaFactoryBeanPostProcessorTests.java @@ -0,0 +1,127 @@ +/* + * 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.util.ArrayList; +import java.util.List; + +import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.producer.Producer; +import org.junit.jupiter.api.Test; + +import org.springframework.kafka.core.ConsumerFactory; +import org.springframework.kafka.core.ConsumerPostProcessor; +import org.springframework.kafka.core.ProducerFactory; +import org.springframework.kafka.core.ProducerPostProcessor; + +import static org.assertj.core.api.BDDAssertions.then; + +class KafkaFactoryBeanPostProcessorTests { + + @Test + void should_add_consumer_post_processor_when_one_is_not_present() { + TestConsumerFactory factory = new TestConsumerFactory(); + then(factory.postProcessors).isEmpty(); + KafkaFactoryBeanPostProcessor processor = new KafkaFactoryBeanPostProcessor(null); + + processor.postProcessAfterInitialization(factory, ""); + + then(factory.postProcessors).hasSize(1); + } + + @Test + void should_add_producer_post_processor_when_one_is_not_present() { + TestProducerFactory factory = new TestProducerFactory(); + then(factory.postProcessors).isEmpty(); + KafkaFactoryBeanPostProcessor processor = new KafkaFactoryBeanPostProcessor(null); + + processor.postProcessAfterInitialization(factory, ""); + + then(factory.postProcessors).hasSize(1); + } + + @Test + void should_not_add_consumer_post_processor_when_one_is_present() { + TestConsumerFactory factory = new TestConsumerFactory(); + factory.postProcessors.add(new TraceConsumerPostProcessor(null)); + KafkaFactoryBeanPostProcessor processor = new KafkaFactoryBeanPostProcessor(null); + + processor.postProcessAfterInitialization(factory, ""); + + then(factory.postProcessors).hasSize(1); + } + + @Test + void should_not_add_producer_post_processor_when_one_is_present() { + TestProducerFactory factory = new TestProducerFactory(); + factory.postProcessors.add(new TraceProducerPostProcessor(null)); + KafkaFactoryBeanPostProcessor processor = new KafkaFactoryBeanPostProcessor(null); + + processor.postProcessAfterInitialization(factory, ""); + + then(factory.postProcessors).hasSize(1); + } + +} + +class TestConsumerFactory implements ConsumerFactory { + + List postProcessors = new ArrayList<>(); + + @Override + public Consumer createConsumer(String groupId, String clientIdPrefix, + String clientIdSuffix) { + return null; + } + + @Override + public boolean isAutoCommit() { + return false; + } + + @Override + public void addPostProcessor(ConsumerPostProcessor postProcessor) { + this.postProcessors.add(postProcessor); + } + + @Override + public List getPostProcessors() { + return this.postProcessors; + } + +} + +class TestProducerFactory implements ProducerFactory { + + List postProcessors = new ArrayList<>(); + + @Override + public Producer createProducer() { + return null; + } + + @Override + public void addPostProcessor(ProducerPostProcessor postProcessor) { + this.postProcessors.add(postProcessor); + } + + @Override + public List getPostProcessors() { + return this.postProcessors; + } + +}