From 013d95120a7b54b3d75ec4daaab547051ecdcf8c Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Fri, 30 Jul 2021 11:38:07 +0200 Subject: [PATCH] Added an option to mutate the reactor context; fixes gh-2000 --- .../brave/BraveBridgeConfiguration.java | 43 +++++++++++++ .../bridge/BraveContextWrappingFunction.java | 60 ++++++++++++++++++ .../cloud/sleuth/brave/bridge/BraveSpan.java | 6 +- .../BraveContextWrappingFunctionTests.java | 62 +++++++++++++++++++ ...eactorSleuthMethodInvocationProcessor.java | 4 +- .../instrument/reactor/ReactorSleuth.java | 19 +++++- .../reactor/ScopePassingSpanSubscriber.java | 3 +- .../sleuth/instrument/web/TraceWebFilter.java | 3 +- .../client/HttpClientBeanPostProcessor.java | 3 +- .../client/TraceExchangeFilterFunction.java | 4 +- 10 files changed, 199 insertions(+), 8 deletions(-) create mode 100644 spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunction.java create mode 100644 spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunctionTests.java diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/brave/BraveBridgeConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/brave/BraveBridgeConfiguration.java index 81626fe1b..cc446473e 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/brave/BraveBridgeConfiguration.java +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/brave/BraveBridgeConfiguration.java @@ -16,19 +16,30 @@ package org.springframework.cloud.sleuth.autoconfig.brave; +import java.io.Closeable; +import java.io.IOException; import java.util.ArrayList; import java.util.List; +import java.util.function.Function; import brave.Tracing; import brave.handler.SpanHandler; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.springframework.beans.BeansException; import org.springframework.beans.factory.BeanFactory; import org.springframework.beans.factory.ObjectProvider; +import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; +import org.springframework.beans.factory.support.BeanDefinitionRegistry; +import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; import org.springframework.boot.context.properties.EnableConfigurationProperties; import org.springframework.cloud.sleuth.SpanCustomizer; import org.springframework.cloud.sleuth.autoconfig.SleuthBaggageProperties; import org.springframework.cloud.sleuth.brave.bridge.BraveBaggageManager; +import org.springframework.cloud.sleuth.brave.bridge.BraveContextWrappingFunction; import org.springframework.cloud.sleuth.brave.bridge.BraveCurrentTraceContext; import org.springframework.cloud.sleuth.brave.bridge.BravePropagator; import org.springframework.cloud.sleuth.brave.bridge.BraveSpanCustomizer; @@ -38,6 +49,7 @@ import org.springframework.cloud.sleuth.brave.bridge.CompositeSpanHandler; import org.springframework.cloud.sleuth.brave.propagation.PropagationFactorySupplier; import org.springframework.cloud.sleuth.exporter.SpanFilter; import org.springframework.cloud.sleuth.exporter.SpanReporter; +import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth; import org.springframework.cloud.sleuth.propagation.Propagator; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -83,4 +95,35 @@ class BraveBridgeConfiguration { reporters.getIfAvailable(ArrayList::new)); } + @Bean + @ConditionalOnClass(name = "reactor.util.context.Context") + static BraveReactorContextBeanDefinitionRegistryPostProcessor braveReactorContextBeanDefinitionRegistryPostProcessor() { + return new BraveReactorContextBeanDefinitionRegistryPostProcessor(); + } + + static class BraveReactorContextBeanDefinitionRegistryPostProcessor + implements BeanDefinitionRegistryPostProcessor, Closeable { + + private static final Log log = LogFactory.getLog(BraveReactorContextBeanDefinitionRegistryPostProcessor.class); + + @Override + public void close() throws IOException { + ReactorSleuth.contextWrappingFunction = Function.identity(); + } + + @Override + public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) throws BeansException { + ReactorSleuth.contextWrappingFunction = new BraveContextWrappingFunction(); + if (log.isDebugEnabled()) { + log.debug("Wrapped Reactor's context into a Brave representation"); + } + } + + @Override + public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException { + + } + + } + } diff --git a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunction.java b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunction.java new file mode 100644 index 000000000..13cf9d942 --- /dev/null +++ b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunction.java @@ -0,0 +1,60 @@ +/* + * 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.function.Function; + +import reactor.util.context.Context; + +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.TraceContext; + +/** + * A function that wraps {@link Context} with Brave versions. + * + * @author Marcin Grzejszczak + * @since 3.0.4 + */ +public class BraveContextWrappingFunction implements Function { + + @Override + public Context apply(Context context) { + Span span = context.getOrDefault(Span.class, null); + TraceContext traceContext = context.getOrDefault(TraceContext.class, null); + if (span == null && traceContext == null) { + return context; + } + if (context.hasKey(brave.propagation.TraceContext.class)) { + return context; + } + return mutateContextWithBrave(context, span, traceContext); + } + + private Context mutateContextWithBrave(Context context, Span span, TraceContext traceContext) { + brave.Span braveSpan = BraveSpan.toBrave(span); + brave.propagation.TraceContext braveTraceContext = BraveTraceContext.toBrave(traceContext); + Context mutatedContext = context; + if (braveSpan != null) { + mutatedContext = context.put(brave.Span.class, braveSpan); + } + if (braveTraceContext != null) { + mutatedContext = mutatedContext.put(brave.propagation.TraceContext.class, braveTraceContext); + } + return mutatedContext; + } + +} diff --git a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveSpan.java b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveSpan.java index 98ea9bd9d..ee4e661e0 100644 --- a/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveSpan.java +++ b/spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveSpan.java @@ -109,7 +109,11 @@ public class BraveSpan implements Span { } public static brave.Span toBrave(Span span) { - return ((BraveSpan) AssertingSpan.unwrap(span)).delegate; + BraveSpan unwrap = (BraveSpan) AssertingSpan.unwrap(span); + if (unwrap == null) { + return null; + } + return unwrap.delegate; } public static Span fromBrave(brave.Span span) { diff --git a/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunctionTests.java b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunctionTests.java new file mode 100644 index 000000000..9b2707351 --- /dev/null +++ b/spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunctionTests.java @@ -0,0 +1,62 @@ +/* + * 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 org.junit.jupiter.api.Test; +import reactor.util.context.Context; + +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.TraceContext; + +import static org.assertj.core.api.BDDAssertions.then; +import static org.mockito.Mockito.mock; + +class BraveContextWrappingFunctionTests { + + BraveContextWrappingFunction function = new BraveContextWrappingFunction(); + + @Test + void should_not_mutate_context_when_no_tracing_information_is_set() { + Context context = Context.empty(); + + then(this.function.apply(context)).isSameAs(context); + } + + @Test + void should_not_mutate_context_when_brave_trace_context_is_already_there() { + Context context = Context.of(Span.class, mock(Span.class), TraceContext.class, mock(TraceContext.class), + brave.propagation.TraceContext.class, traceContext()); + + then(this.function.apply(context)).isSameAs(context); + } + + @Test + void should_mutate_context_when_there_is_tracing_info_but_brave_version_is_missing() { + Context context = Context.of(Span.class, new BraveSpan(mock(brave.Span.class)), TraceContext.class, + new BraveTraceContext(traceContext())); + + Context mutatedContext = this.function.apply(context); + + then(mutatedContext.hasKey(brave.Span.class)).isTrue(); + then(mutatedContext.hasKey(brave.propagation.TraceContext.class)).isTrue(); + } + + private brave.propagation.TraceContext traceContext() { + return brave.propagation.TraceContext.newBuilder().spanId(1L).traceId(2L).build(); + } + +} diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java index 5191b6a04..3114ff539 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java @@ -36,6 +36,7 @@ import org.springframework.cloud.sleuth.TraceContext; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.annotation.ContinueSpan; import org.springframework.cloud.sleuth.annotation.NewSpan; +import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth; import org.springframework.cloud.sleuth.instrument.reactor.TraceContextPropagator; import org.springframework.util.StringUtils; @@ -228,7 +229,8 @@ public class ReactorSleuthMethodInvocationProcessor extends AbstractSleuthMethod this.log = log; this.hasLog = hasLog; this.processor = processor; - this.context = actual.currentContext().put(Span.class, span).put(TraceContext.class, span.context()); + this.context = ReactorSleuth + .wrapContext(actual.currentContext().put(Span.class, span).put(TraceContext.class, span.context())); this.tracer = processor.tracer(); processor.before(invocation, this.span, this.log, this.hasLog); } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java index 3c2063349..2c4ecff5c 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java @@ -64,6 +64,11 @@ public abstract class ReactorSleuth { private ReactorSleuth() { } + /** + * Function that does additional wrapping of the Reactor context. + */ + public static Function contextWrappingFunction = Function.identity(); + /** * Return a span operator pointcut given a Tracing. This can be used in reactor via * {@link reactor.core.publisher.Flux#transform(Function)}, @@ -582,8 +587,18 @@ public abstract class ReactorSleuth { * @return mutated context */ public static Context putSpanInScope(Tracer tracer, Context context, Span span) { - return context.put(Span.class, span).put(TraceContext.class, span.context()).put(Tracer.SpanInScope.class, - tracer.withSpan(span)); + Context newContext = context.put(Span.class, span).put(TraceContext.class, span.context()) + .put(Tracer.SpanInScope.class, tracer.withSpan(span)); + return wrapContext(newContext); + } + + /** + * Mutates the Reactor context depending on the classpath contents. + * @param context Reactor context + * @return mutated context + */ + public static Context wrapContext(Context context) { + return contextWrappingFunction.apply(context); } /** diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java index 8c62ec0a7..907788471 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java @@ -53,8 +53,9 @@ final class ScopePassingSpanSubscriber implements SpanSubscription, Scanna this.subscriber = subscriber; this.currentTraceContext = currentTraceContext; this.parent = parent; - this.context = parent != null && !parent.equals(ctx.getOrDefault(TraceContext.class, null)) + Context context = parent != null && !parent.equals(ctx.getOrDefault(TraceContext.class, null)) ? ctx.put(TraceContext.class, parent) : ctx; + this.context = ReactorSleuth.wrapContext(context); if (log.isTraceEnabled()) { log.trace("Parent span [" + parent + "], context [" + this.context + "]"); } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java index df64b85b7..8237da9f2 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java @@ -37,6 +37,7 @@ import org.springframework.cloud.sleuth.docs.AssertingSpan; import org.springframework.cloud.sleuth.http.HttpServerHandler; import org.springframework.cloud.sleuth.http.HttpServerRequest; import org.springframework.cloud.sleuth.http.HttpServerResponse; +import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth; import org.springframework.cloud.sleuth.instrument.reactor.TraceContextPropagator; import org.springframework.context.ApplicationContext; import org.springframework.context.ApplicationContextAware; @@ -253,7 +254,7 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa MonoWebFilterTrace parent) { this.actual = actual; this.span = span; - this.context = context.put(TraceContext.class, span.context()); + this.context = ReactorSleuth.wrapContext(context.put(TraceContext.class, span.context())); this.exchange = parent.exchange; this.handler = parent.handler; } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java index 4f819f0ec..0f2258c4f 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java @@ -40,6 +40,7 @@ import org.springframework.cloud.sleuth.CurrentTraceContext; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.TraceContext; import org.springframework.cloud.sleuth.http.HttpClientHandler; +import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth; import org.springframework.cloud.sleuth.internal.LazyBean; import org.springframework.context.ConfigurableApplicationContext; import org.springframework.lang.Nullable; @@ -112,7 +113,7 @@ public class HttpClientBeanPostProcessor implements BeanPostProcessor { TraceContext invocationContext = currentTraceContext.get(); if (invocationContext != null) { // Read in this processor and also in ScopePassingSpanSubscriber - context = context.put(TraceContext.class, invocationContext); + context = ReactorSleuth.wrapContext(context.put(TraceContext.class, invocationContext)); } return context.put(PendingSpan.class, pendingSpan); }).doOnCancel(() -> { diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceExchangeFilterFunction.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceExchangeFilterFunction.java index 9efc62798..3d960236f 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceExchangeFilterFunction.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceExchangeFilterFunction.java @@ -36,6 +36,7 @@ import org.springframework.cloud.sleuth.TraceContext; import org.springframework.cloud.sleuth.http.HttpClientHandler; import org.springframework.cloud.sleuth.http.HttpClientRequest; import org.springframework.cloud.sleuth.http.HttpClientResponse; +import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth; import org.springframework.cloud.sleuth.instrument.reactor.TraceContextPropagator; import org.springframework.context.ConfigurableApplicationContext; import org.springframework.http.HttpMethod; @@ -171,8 +172,9 @@ public final class TraceExchangeFilterFunction implements ExchangeFilterFunction this.currentTraceContext = mono.currentTraceContext; this.method = mono.request.method(); this.httpRoute = (String) mono.request.attribute(URI_TEMPLATE_ATTRIBUTE).orElse(null); - this.context = this.parent != null && !this.parent.equals(ctx.getOrDefault(TraceContext.class, null)) + Context context = this.parent != null && !this.parent.equals(ctx.getOrDefault(TraceContext.class, null)) ? ctx.put(TraceContext.class, this.parent) : ctx; + this.context = ReactorSleuth.wrapContext(context); set(clientSpan); }