SpringCloud/spring-cloud-sleuth
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Added an option to mutate the reactor context; fixes gh-2000

Browse Source
This commit is contained in:
Marcin Grzejszczak
2021-07-30 11:38:07 +02:00
parent 69691471ad
commit 013d95120a
10 changed files with 199 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/brave/BraveBridgeConfiguration.java
View File

@@ -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 {
}
}
}

60
spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunction.java Normal file
View File

@@ -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<Context, Context> {
@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;
}
}

6
spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveSpan.java
View File

@@ -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) {

62
spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/BraveContextWrappingFunctionTests.java Normal file
View File

@@ -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();
}
}

4
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/annotation/ReactorSleuthMethodInvocationProcessor.java
View File

@@ -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);
}

19
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java
View File

@@ -64,6 +64,11 @@ public abstract class ReactorSleuth {
private ReactorSleuth() {
}
/**
* Function that does additional wrapping of the Reactor context.
*/
public static Function<Context, Context> 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);
}
/**

3
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java
View File

@@ -53,8 +53,9 @@ final class ScopePassingSpanSubscriber<T> implements SpanSubscription<T>, 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 + "]");
}

3
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java
View File

@@ -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;
}

3
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java
View File

@@ -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(() -> {

4
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceExchangeFilterFunction.java
View File

@@ -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);
}