Wraps producer / consumer once only (#1923)

* Wraps producer / consumer once only

without this change we can wrap the kafka consumer / producer multiple times. First via post processors, second via aspect
with this change we're doing it only once via post processors. we do not wrap factories

fixes gh-1921

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/instrument/messaging/SleuthKafkaAspect.java

@@ -22,8 +22,6 @@ import brave.Tracer;
 import brave.kafka.clients.KafkaTracing;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.clients.producer.Producer;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
@@ -60,18 +58,6 @@ public class SleuthKafkaAspect {
 				"recordMessageConverter");
 	}
 
-	@Pointcut("execution(public * org.springframework.kafka.core.ProducerFactory.createProducer(..))")
-	private void anyProducerFactory() {
-	} // NOSONAR
-
-	@Pointcut("execution(public * org.springframework.kafka.core.ProducerFactory.createNonTransactionalProducer(..))")
-	private void anyNonTransactionalProducerFactory() {
-	} // NOSONAR
-
-	@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
@@ -80,18 +66,6 @@ public class SleuthKafkaAspect {
 	private void anyCreateContainer() {
 	} // NOSONAR
 
-	@Around("anyProducerFactory() || anyNonTransactionalProducerFactory()")
-	public Object wrapProducerFactory(ProceedingJoinPoint pjp) throws Throwable {
-		Producer producer = (Producer) pjp.proceed();
-		return this.kafkaTracing.producer(producer);
-	}
-
-	@Around("anyConsumerFactory()")
-	public Object wrapConsumerFactory(ProceedingJoinPoint pjp) throws Throwable {
-		Consumer consumer = (Consumer) pjp.proceed();
-		return this.kafkaTracing.consumer(consumer);
-	}
-
 	@Around("anyCreateListenerContainer() || anyCreateContainer()")
 	public Object wrapListenerContainerCreation(ProceedingJoinPoint pjp) throws Throwable {
 		MessageListenerContainer listener = (MessageListenerContainer) pjp.proceed();

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/instrument/messaging/TraceConsumerPostProcessor.java

@@ -21,6 +21,7 @@ import org.apache.kafka.clients.consumer.Consumer;
 
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.kafka.core.ConsumerPostProcessor;
+import org.springframework.util.ClassUtils;
 
 class TraceConsumerPostProcessor<K, V> implements ConsumerPostProcessor<K, V> {
 
@@ -28,6 +29,10 @@ class TraceConsumerPostProcessor<K, V> implements ConsumerPostProcessor<K, V> {
 
 	private KafkaTracing kafkaTracing;
 
+	// Because it's not public in Brave
+	private static final Class tracingConsumer = ClassUtils.resolveClassName("brave.kafka.clients.TracingConsumer",
+			null);
+
 	TraceConsumerPostProcessor(BeanFactory beanFactory) {
 		this.beanFactory = beanFactory;
 	}
@@ -41,6 +46,13 @@ class TraceConsumerPostProcessor<K, V> implements ConsumerPostProcessor<K, V> {
 
 	@Override
 	public Consumer<K, V> apply(Consumer<K, V> kvConsumer) {
+		if (tracingConsumer.isAssignableFrom(ClassUtils.getUserClass(kvConsumer.getClass()))) {
+			return kvConsumer;
+		}
+		return wrapInTracing(kvConsumer);
+	}
+
+	Consumer<K, V> wrapInTracing(Consumer<K, V> kvConsumer) {
 		return kafkaTracing().consumer(kvConsumer);
 	}
 

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/instrument/messaging/TraceProducerPostProcessor.java

@@ -21,6 +21,7 @@ import org.apache.kafka.clients.producer.Producer;
 
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.kafka.core.ProducerPostProcessor;
+import org.springframework.util.ClassUtils;
 
 class TraceProducerPostProcessor<K, V> implements ProducerPostProcessor<K, V> {
 
@@ -28,6 +29,10 @@ class TraceProducerPostProcessor<K, V> implements ProducerPostProcessor<K, V> {
 
 	private KafkaTracing kafkaTracing;
 
+	// Because it's not public in Brave
+	private static final Class tracingProducer = ClassUtils.resolveClassName("brave.kafka.clients.TracingProducer",
+			null);
+
 	TraceProducerPostProcessor(BeanFactory beanFactory) {
 		this.beanFactory = beanFactory;
 	}
@@ -41,6 +46,13 @@ class TraceProducerPostProcessor<K, V> implements ProducerPostProcessor<K, V> {
 
 	@Override
 	public Producer<K, V> apply(Producer<K, V> kvProducer) {
+		if (tracingProducer.isAssignableFrom(ClassUtils.getUserClass(kvProducer.getClass()))) {
+			return kvProducer;
+		}
+		return wrapInTracing(kvProducer);
+	}
+
+	Producer<K, V> wrapInTracing(Producer<K, V> kvProducer) {
 		return kafkaTracing().producer(kvProducer);
 	}
 

spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/instrument/messaging/TraceConsumerPostProcessorTests.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2013-2021 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.brave.instrument.messaging;
+
+import brave.Tracing;
+import brave.kafka.clients.KafkaTracing;
+import brave.propagation.StrictCurrentTraceContext;
+import brave.sampler.Sampler;
+import brave.test.TestSpanHandler;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.assertj.core.api.BDDAssertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.BDDMockito;
+
+import org.springframework.beans.factory.BeanFactory;
+
+class TraceConsumerPostProcessorTests {
+
+	TestSpanHandler spans = new TestSpanHandler();
+
+	StrictCurrentTraceContext traceContext = StrictCurrentTraceContext.create();
+
+	Tracing tracing = Tracing.newBuilder().currentTraceContext(this.traceContext).sampler(Sampler.ALWAYS_SAMPLE)
+			.addSpanHandler(this.spans).build();
+
+	KafkaTracing kafkaTracing = KafkaTracing.newBuilder(this.tracing).build();
+
+	@Test
+	void should_not_double_wrap_kafka_consumer() {
+		BeanFactory beanFactory = BDDMockito.mock(BeanFactory.class);
+		BDDMockito.given(beanFactory.getBean(KafkaTracing.class)).willReturn(kafkaTracing);
+		Consumer consumer = BDDMockito.mock(Consumer.class);
+		Consumer wrappedConsumer = kafkaTracing.consumer(consumer);
+
+		final Consumer apply = new TraceConsumerPostProcessor(beanFactory) {
+			@Override
+			Consumer wrapInTracing(Consumer consumer) {
+				throw new AssertionError("This method must not be called");
+			}
+		}.apply(wrappedConsumer);
+
+		BDDAssertions.then(apply).isSameAs(wrappedConsumer);
+	}
+
+}

spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/instrument/messaging/TraceProducerPostProcessorTests.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2013-2021 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.brave.instrument.messaging;
+
+import brave.Tracing;
+import brave.kafka.clients.KafkaTracing;
+import brave.propagation.StrictCurrentTraceContext;
+import brave.sampler.Sampler;
+import brave.test.TestSpanHandler;
+import org.apache.kafka.clients.producer.Producer;
+import org.assertj.core.api.BDDAssertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.BDDMockito;
+
+import org.springframework.beans.factory.BeanFactory;
+
+class TraceProducerPostProcessorTests {
+
+	TestSpanHandler spans = new TestSpanHandler();
+
+	StrictCurrentTraceContext traceContext = StrictCurrentTraceContext.create();
+
+	Tracing tracing = Tracing.newBuilder().currentTraceContext(this.traceContext).sampler(Sampler.ALWAYS_SAMPLE)
+			.addSpanHandler(this.spans).build();
+
+	KafkaTracing kafkaTracing = KafkaTracing.newBuilder(this.tracing).build();
+
+	@Test
+	void should_not_double_wrap_kafka_producer() {
+		BeanFactory beanFactory = BDDMockito.mock(BeanFactory.class);
+		BDDMockito.given(beanFactory.getBean(KafkaTracing.class)).willReturn(kafkaTracing);
+		Producer producer = BDDMockito.mock(Producer.class);
+		Producer wrappedProducer = kafkaTracing.producer(producer);
+
+		final Producer apply = new TraceProducerPostProcessor(beanFactory) {
+			@Override
+			Producer wrapInTracing(Producer producer) {
+				throw new AssertionError("This method must not be called");
+			}
+		}.apply(wrappedProducer);
+
+		BDDAssertions.then(apply).isSameAs(wrappedProducer);
+	}
+
+}

tests/brave/spring-cloud-sleuth-instrumentation-messaging-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/messaging/BraveMessagingAutoConfigurationTests.java

@@ -26,9 +26,7 @@ import brave.sampler.SamplerFunction;
 import brave.sampler.SamplerFunctions;
 import brave.spring.rabbit.SpringRabbitTracing;
 import brave.test.TestSpanHandler;
-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.junit.jupiter.api.Test;
 import org.mockito.Mockito;
@@ -94,16 +92,6 @@ public class BraveMessagingAutoConfigurationTests {
 
 	@Test
 	public void should_wrap_kafka() {
-		this.producerFactory.createProducer();
-		then(this.mySleuthKafkaAspect.producerWrapped).isTrue();
-		this.mySleuthKafkaAspect.producerWrapped = false;
-
-		this.producerFactory.createNonTransactionalProducer();
-		then(this.mySleuthKafkaAspect.producerWrapped).isTrue();
-
-		this.consumerFactory.createConsumer();
-		then(this.mySleuthKafkaAspect.consumerWrapped).isTrue();
-
 		then(this.mySleuthKafkaAspect.adapterWrapped).isTrue();
 	}
 
@@ -208,28 +196,12 @@ class TestSleuthRabbitBeanPostProcessor extends SleuthRabbitBeanPostProcessor {
 
 class MySleuthKafkaAspect extends SleuthKafkaAspect {
 
-	boolean producerWrapped;
-
-	boolean consumerWrapped;
-
 	boolean adapterWrapped;
 
 	MySleuthKafkaAspect(KafkaTracing kafkaTracing, Tracer tracer) {
 		super(kafkaTracing, tracer);
 	}
 
-	@Override
-	public Object wrapProducerFactory(ProceedingJoinPoint pjp) throws Throwable {
-		this.producerWrapped = true;
-		return Mockito.mock(Producer.class);
-	}
-
-	@Override
-	public Object wrapConsumerFactory(ProceedingJoinPoint pjp) throws Throwable {
-		this.consumerWrapped = true;
-		return Mockito.mock(Consumer.class);
-	}
-
 	@Override
 	public Object wrapListenerContainerCreation(ProceedingJoinPoint pjp) throws Throwable {
 		this.adapterWrapped = true;