From 984f62c32d8278e9bbd31b6c9698c6b533d42c97 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Tue, 16 Mar 2021 17:00:56 +0100 Subject: [PATCH] Fixed custom propagation mode option with this change when CUSTOM mode is turned on we will search for user provided Propagation bean or else noop fixes gh-1836 --- .../CompositePropagationFactorySupplier.java | 35 +++++- ...positePropagationFactorySupplierTests.java | 118 ++++++++++++++++++ .../TraceWebClientBeanPostProcessor.java | 1 - .../TraceFeignBuilderBeanPostProcessor.java | 4 +- 4 files changed, 151 insertions(+), 7 deletions(-) create mode 100644 spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/CompositePropagationFactorySupplierTests.java diff --git a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/CompositePropagationFactorySupplier.java b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/CompositePropagationFactorySupplier.java index 787d02fc2..167ccce25 100644 --- a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/CompositePropagationFactorySupplier.java +++ b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/CompositePropagationFactorySupplier.java @@ -30,6 +30,7 @@ import brave.propagation.TraceContextOrSamplingFlags; import brave.propagation.aws.AWSPropagation; import org.springframework.beans.factory.BeanFactory; +import org.springframework.beans.factory.ObjectProvider; import org.springframework.cloud.sleuth.brave.propagation.PropagationFactorySupplier; import org.springframework.cloud.sleuth.brave.propagation.PropagationType; @@ -56,7 +57,7 @@ public class CompositePropagationFactorySupplier implements PropagationFactorySu @Override public Propagation.Factory get() { - return new CompositePropagationFactory( + return new CompositePropagationFactory(this.beanFactory, this.beanFactory.getBeanProvider(BraveBaggageManager.class).getIfAvailable(BraveBaggageManager::new), this.localFields, this.types); } @@ -69,8 +70,8 @@ class CompositePropagationFactory extends Propagation.Factory implements Propaga private final List types; - CompositePropagationFactory(BraveBaggageManager braveBaggageManager, List localFields, - List types) { + CompositePropagationFactory(BeanFactory beanFactory, BraveBaggageManager braveBaggageManager, + List localFields, List types) { this.types = types; this.mapping.put(PropagationType.AWS, AWSPropagation.FACTORY.get()); // Note: Versions <2.2.3 use injectFormat(MULTI) for non-remote (ex @@ -79,7 +80,7 @@ class CompositePropagationFactory extends Propagation.Factory implements Propaga this.mapping.put(PropagationType.B3, B3Propagation.newFactoryBuilder().injectFormat(B3Propagation.Format.SINGLE_NO_PARENT).build().get()); this.mapping.put(PropagationType.W3C, new W3CPropagation(braveBaggageManager, localFields)); - this.mapping.put(PropagationType.CUSTOM, NoOpPropagation.INSTANCE); + this.mapping.put(PropagationType.CUSTOM, new LazyPropagation(beanFactory.getBeanProvider(Propagation.class))); } @Override @@ -116,6 +117,32 @@ class CompositePropagationFactory extends Propagation.Factory implements Propaga return StringPropagationAdapter.create(this, keyFactory); } + @SuppressWarnings("unchecked") + private static final class LazyPropagation implements Propagation { + + private final ObjectProvider delegate; + + private LazyPropagation(ObjectProvider delegate) { + this.delegate = delegate; + } + + @Override + public List keys() { + return this.delegate.getIfAvailable(() -> NoOpPropagation.INSTANCE).keys(); + } + + @Override + public TraceContext.Injector injector(Setter setter) { + return this.delegate.getIfAvailable(() -> NoOpPropagation.INSTANCE).injector(setter); + } + + @Override + public TraceContext.Extractor extractor(Getter getter) { + return this.delegate.getIfAvailable(() -> NoOpPropagation.INSTANCE).extractor(getter); + } + + } + private static class NoOpPropagation implements Propagation { static final NoOpPropagation INSTANCE = new NoOpPropagation(); diff --git a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/CompositePropagationFactorySupplierTests.java b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/CompositePropagationFactorySupplierTests.java new file mode 100644 index 000000000..1884a69a1 --- /dev/null +++ b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/CompositePropagationFactorySupplierTests.java @@ -0,0 +1,118 @@ +/* + * 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.bridge; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +import brave.internal.codec.HexCodec; +import brave.internal.propagation.StringPropagationAdapter; +import brave.propagation.Propagation; +import brave.propagation.TraceContext; +import brave.propagation.TraceContextOrSamplingFlags; +import org.assertj.core.api.BDDAssertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import org.springframework.beans.factory.BeanFactory; +import org.springframework.cloud.loadbalancer.support.SimpleObjectProvider; +import org.springframework.cloud.sleuth.brave.propagation.PropagationType; +import org.springframework.util.StringUtils; + +class CompositePropagationFactorySupplierTests { + + @Test + void should_pick_custom_registered_propagation_when_custom_mode_picked() { + BeanFactory beanFactory = Mockito.mock(BeanFactory.class); + Mockito.when(beanFactory.getBeanProvider(BraveBaggageManager.class)) + .thenReturn(new SimpleObjectProvider(new BraveBaggageManager())); + Mockito.when(beanFactory.getBeanProvider(Propagation.class)) + .thenReturn(new SimpleObjectProvider(new CustomTracePropagation())); + + CompositePropagationFactorySupplier supplier = new CompositePropagationFactorySupplier(beanFactory, + Collections.emptyList(), Collections.singletonList(PropagationType.CUSTOM)); + + BDDAssertions.then(supplier.get().get().keys()).containsExactly(CustomTraceExtractor.CUSTOM_TRACE_HEADER); + } + +} + +class CustomTracePropagation extends Propagation.Factory implements Propagation { + + public static final List KEYS = Collections.singletonList(CustomTraceExtractor.CUSTOM_TRACE_HEADER); + + @Override + public List keys() { + return KEYS; + } + + @Override + public TraceContext.Injector injector(Setter setter) { + return (traceContext, request) -> { + String trace = traceContext.traceIdString() + ":" + traceContext.spanIdString(); + setter.put(request, CustomTraceExtractor.CUSTOM_TRACE_HEADER, trace); + }; + } + + @Override + public TraceContext.Extractor extractor(Getter getter) { + Objects.requireNonNull(getter); + return new CustomTraceExtractor<>(getter); + } + + @Override + public Propagation create(KeyFactory keyFactory) { + return StringPropagationAdapter.create(this, keyFactory); + } + +} + +class CustomTraceExtractor implements TraceContext.Extractor { + + static final String CUSTOM_TRACE_HEADER = "x-custom-trace"; + + final Propagation.Getter getter; + + CustomTraceExtractor(Propagation.Getter getter) { + this.getter = getter; + } + + @Override + @SuppressWarnings("ReturnCount") + public TraceContextOrSamplingFlags extract(R request) { + String traceString = getter.get(request, CUSTOM_TRACE_HEADER); + if (!StringUtils.hasText(traceString)) { + return TraceContextOrSamplingFlags.EMPTY; + } + String[] trace = traceString.split(":"); + if (trace.length != 2) { + return TraceContextOrSamplingFlags.EMPTY; + } + + try { + TraceContext traceContext = TraceContext.newBuilder().traceId(HexCodec.lowerHexToUnsignedLong(trace[0])) + .spanId(HexCodec.lowerHexToUnsignedLong(trace[1])).build(); + + return TraceContextOrSamplingFlags.create(traceContext); + } + catch (NumberFormatException ex) { + return TraceContextOrSamplingFlags.EMPTY; + } + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java index 6a0127801..1dd21b4df 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java @@ -80,4 +80,3 @@ public class TraceWebClientBeanPostProcessor implements BeanPostProcessor { } } - diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignBuilderBeanPostProcessor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignBuilderBeanPostProcessor.java index 168f0314d..c10596245 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignBuilderBeanPostProcessor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignBuilderBeanPostProcessor.java @@ -27,8 +27,8 @@ import org.springframework.beans.factory.config.BeanPostProcessor; import org.springframework.util.ReflectionUtils; /** - * {@link BeanPostProcessor} that ensures that each {@link Feign.Builder} has - * a trace representation of a {@link Client}. + * {@link BeanPostProcessor} that ensures that each {@link Feign.Builder} has a trace + * representation of a {@link Client}. * * @author Marcin Grzejszczak * @since 3.0.2