Added Consumer / Producer post processors (#1701)

fixes gh-1670

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<K, V> implements ConsumerPostProcessor<K, V> {
+
+	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<K, V> apply(Consumer<K, V> kvConsumer) {
+		return kafkaTracing().consumer(kvConsumer);
+	}
+
+}
+
+class TraceProducerPostProcessor<K, V> implements ProducerPostProcessor<K, V> {
+
+	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<K, V> apply(Producer<K, V> kvProducer) {
+		return kafkaTracing().producer(kvProducer);
+	}
+
+}
+
 @Aspect
 class SleuthKafkaAspect {
 

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<ConsumerPostProcessor> 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<ConsumerPostProcessor> getPostProcessors() {
+		return this.postProcessors;
+	}
+
+}
+
+class TestProducerFactory implements ProducerFactory {
+
+	List<ProducerPostProcessor> postProcessors = new ArrayList<>();
+
+	@Override
+	public Producer createProducer() {
+		return null;
+	}
+
+	@Override
+	public void addPostProcessor(ProducerPostProcessor postProcessor) {
+		this.postProcessors.add(postProcessor);
+	}
+
+	@Override
+	public List<ProducerPostProcessor> getPostProcessors() {
+		return this.postProcessors;
+	}
+
+}